Skip to content

CONTRIBUTING.md

Latest commit

 

History

History
154 lines (105 loc) · 8.02 KB

CONTRIBUTING.md

File metadata and controls

154 lines (105 loc) · 8.02 KB
Raw

Contributing To Emerging-Optimizers

Dependency

Use abseil-py's logging, testing and flags instead of Python's own logging, unittest and argparse.

We use uv for managing dependencies. For reproducible builds, our project tracks the generated uv.lock file in the repository. On a weekly basis, the CI attempts an update of the lock file to test against upstream dependencies.

New required dependencies can be added by uv add $DEPENDENCY.

New optional dependencies can be added by uv add --optional --extra $EXTRA $DEPENDENCY.

EXTRA refers to the subgroup of extra-dependencies to which you're adding the new dependency. Example: For adding a TRT-LLM specific dependency, run uv add --optional --extra trtllm $DEPENDENCY.

New dependencies to a dependency group can be added with uv add --group $GROUP $DEPENDENCY. Dependency groups are specific to uv and are also optional dependencies but not intended to be used with the package such as docs dependencies.

Alternatively, the pyproject.toml file can also be modified directly.

Adding a new dependency will update UV's lock-file. Please check this into your branch:

git add uv.lock pyproject.toml
git commit -m "build: Adding dependencies"
git push

🧹 Linting and Formatting

We use ruff for linting and formatting. CI does not auto-fix linting and formatting issues, but most issues can be fixed by running the following command:

uv run ruff check --fix .
uv run ruff format .

Coding Style

We generally follow Google's style guides , with some exceptions:

  • Line length extended to 120 for Python and 100 for C++ code.
  • Common use in PyTorch, which are prohibited by Google style, are allowed, including but not limited to:
    • Import function, class, not just module
    • Some special variable name, x, dX, etc.
  • Allow common capitalized naming in Triton code.

Although common, mixed case is not allowed in any code. "Mixed case" here means camelCase (e.g. getFoo, myVariable); PascalCase for classes and snake_case for functions/variables follow Google style as usual.

Run pre-commit at local before submitting merge request. You can also read .pre-commit-config.yaml to understand what are being forced. The mypy settings are inherited from PyTorch.

Optimizer API

closure is not supported on any optimizer's step in this repo. The parameter is kept on the signature for compatibility with torch.optim.Optimizer, but every step implementation raises ValueError("closure is not supported") when one is passed and returns None otherwise. New optimizers must do the same — do not add code paths that call closure() or return its result.

Argument-order convention

When designing or modifying any public function in emerging_optimizers/ — optimizer __init__, calculate_*_update, helpers — follow this precedence for argument order:

  1. PyTorch native first. Match torch.optim / torch.optim._functional signatures. For Adam-family signatures that means betas and eps adjacent (Adam(lr, betas, eps, weight_decay, ...)).
  2. HuggingFace next. For knobs PyTorch doesn't expose, follow HF's order. The clearest case is correct_bias: HF's transformers.AdamW puts it after the standard Adam scalars (lr, betas, eps, weight_decay), so in our update functions correct_bias goes after eps.
  3. Repo-specific extras last. Anything not in PyTorch or HF — nesterov on Adam, the scalar-int step, warmup schedulers, alpha, scale_log2, use_shape_scaling, etc. — goes at the end.

Concrete shape for the scalar calculate_*_update family:

def calculate_*_update(
    grad,                   # positional: tensors only
    <state buffers>,
    *,                      # everything below is keyword-only
    betas (or momentum),
    eps,
    correct_bias,
    [nesterov],
    step,
    <other repo-specific extras>,
):

step is a repo-specific argument and goes in the trailing bucket. (PyTorch's functional adam takes state_steps as a List[Tensor] grouped with the buffers — that is not the same as our scalar step: int, so the PyTorch precedent does not transfer.)

The * keyword-only separator matches the convention used by every optimizer __init__ in emerging_optimizers/ (Muon, Soap, PSGDPro, AdaptiveMuon, etc.) and by torch.optim._functional.adam — only tensors are positional, scalars and flags must be passed by name. mypy will reject positional scalar passes, so the rule is enforced at type-check time.

Test

All tests should be placed under tests. We aim for 100% test coverage for this tiny project.

We use abseil-py testing because it is easier to launch multi process than alternatives.

Every test file must define --device and --seed flags — CI invokes every test under both --device=cpu and --device=cuda, and runs each GPU test twice (random seed, then a fixed seed). A test without these flags will fail at flag parsing in CI. Match the pattern used in existing test files: flags.DEFINE_enum("device", "cpu", ["cpu", "cuda"], ...), flags.DEFINE_integer("seed", None, ...), and a setUpModule that seeds when FLAGS.seed is not None.

When using torch.testing.assert_close, do not override the msg argument with a bare string — that replaces the default diff/atol/rtol summary, which is what makes CI failures debuggable. Pass msg= as a callable that takes the default message and returns one with your context appended. Canonical pattern (from the PyTorch 2.11 docs):

torch.testing.assert_close(
    actual, expected, msg=lambda msg: f"Header\n\n{msg}\n\nFooter"
)

atol and rtol must be set together or not at all — passing only one raises ValueError. Use both or neither (the default tolerances are dtype-aware).

Signing Your Work

  • We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.

    • Any contribution which contains commits that are not Signed-Off will not be accepted.

  • To sign off on a commit you simply use the --signoff (or -s) option when committing your changes:

    $ git commit -s -m "Add cool feature."

    This will append the following to your commit message:

    Signed-off-by: Your Name <your@email.com>

  • Full text of the DCO:

    Developer Certificate of Origin
Version 1.1

Copyright (C) 2004, 2006 The Linux Foundation and its contributors.

Everyone is permitted to copy and distribute verbatim copies of this
license document, but changing it is not allowed.


Developer's Certificate of Origin 1.1

By making a contribution to this project, I certify that:

(a) The contribution was created in whole or in part by me and I
    have the right to submit it under the open source license
    indicated in the file; or

(b) The contribution is based upon previous work that, to the best
    of my knowledge, is covered under an appropriate open source
    license and I have the right under that license to submit that
    work with modifications, whether created in whole or in part
    by me, under the same open source license (unless I am
    permitted to submit under a different license), as indicated
    in the file; or

(c) The contribution was provided directly to me by some other
    person who certified (a), (b) or (c) and I have not modified
    it.

(d) I understand and agree that this project and the contribution
    are public and that a record of the contribution (including all
    personal information I submit with it, including my sign-off) is
    maintained indefinitely and may be redistributed consistent with
    this project or the open source license(s) involved.