with one click
linting-and-formatting
// Code style and quality rules for Megatron Bridge — ruff configuration, naming conventions, type hints, mypy rules, docstrings, copyright headers, logging, and the code review checklist.
// Code style and quality rules for Megatron Bridge — ruff configuration, naming conventions, type hints, mypy rules, docstrings, copyright headers, logging, and the code review checklist.
| name | linting-and-formatting |
| description | Code style and quality rules for Megatron Bridge — ruff configuration, naming conventions, type hints, mypy rules, docstrings, copyright headers, logging, and the code review checklist. |
| when_to_use | Writing or reviewing code for style compliance, fixing ruff or mypy errors, pre-commit hook failures, copyright header questions, naming or docstring conventions. |
Single source of truth for code style in Megatron Bridge. Read this before writing new code or reviewing PRs.
Target Python 3.10+.
Run before every commit:
uv run ruff check --fix .
uv run ruff format .
Pre-commit hooks run these automatically. If hooks auto-fix files, re-stage and re-run until clean:
git add -u
pre-commit run
# if it auto-fixed files:
git add -u
pre-commit run
ruff.toml)| Rule | ID | Description |
|---|---|---|
| Line length | — | 119 characters (formatter) |
| Quote style | — | Double quotes |
| f-string without placeholders | F541 | Error |
| Unused local variable | F841 | Auto-removed by --fix |
| Unused import | F401 | Auto-removed by --fix (ignored in __init__.py) |
| Ambiguous variable name | E741 | Error (e.g., l, O, I) |
| Undefined name | F821 | Error |
| Block comment format | E266 | Error (too many #) |
| Import sorting | I | isort-compatible, auto-fixed |
| Public class docstring | D101 | Warning (ignored in test files) |
| Public function docstring | D103 | Warning (ignored in test files) |
Per-file overrides:
__init__.py: F401 and F403 ignored (re-exports expected).test_*.py, *_test.py, tests/*.py: D101 and D103 ignored.| Kind | Convention | Example |
|---|---|---|
| Files | snake_case | model_bridge.py |
| Classes | PascalCase | MegatronModelBridge |
| Functions/methods | snake_case | load_weights_hf_to_megatron |
| Local variables | snake_case | megatron_weights |
| Variables starting with digit | prefix k | k_99th_percentile |
| Global variables | UPPER_SNAKE + prefix G | G_LOGGER |
| Constants | UPPER_SNAKE | DEFAULT_HIDDEN_SIZE |
__future__ importsmegatron.core, torch, transformers, etc.)megatron.bridge.*)Separate groups with blank lines. ruff auto-fixes via the I rule.
Required on all public API functions and methods.
T | None instead of Optional[T]X | Y instead of Union[X, Y]list, dict, tuple) instead of typing equivalentsdef get_module_by_name(
model: torch.nn.Module,
name: str,
default: torch.nn.Module | None = None,
) -> torch.nn.Module | None:
...
Run on changed files before submitting:
uv run mypy --strict path/to/file.py
Key rules:
Any leaks — use object for unknown types or a TypeVar for generics.Optional — write x: int | None = None, never x: int = None.typing.cast() only when inference fails; add a comment.TypedDict over dict[str, Any] for structured dicts.Callable[[ArgType], ReturnType] or Protocol.# type: ignore[code] must include the error code and justification.When a function has multiple parameters of the same type that could be swapped
by mistake, use * to force keyword-only arguments.
# Don't
def scatter_weights(tensor: Tensor, tp_group: ProcessGroup, ep_group: ProcessGroup): ...
# Do
def scatter_weights(tensor: Tensor, *, tp_group: ProcessGroup, ep_group: ProcessGroup): ...
Google-style for public classes and functions:
def convert_weights(
source_model: torch.nn.Module,
target_model: torch.nn.Module,
mapping: MegatronParamMapping,
) -> dict[str, torch.Tensor]:
"""Convert weights from source to target model format.
Args:
source_model: The source model containing weights to convert.
target_model: The target model that will receive converted weights.
mapping: Parameter mapping defining the conversion rules.
Returns:
Dictionary mapping parameter names to converted weight tensors.
Raises:
ValueError: If source and target models have incompatible shapes.
"""
Use logging.getLogger(__name__) for module-level loggers. Use print_rank_0
/ warn_rank_0 for user-facing messages in distributed contexts.
# Don't
print(f"Loading weights for {model_name}")
# Do
logger = logging.getLogger(__name__)
logger.info("Loading weights for %s", model_name)
Use specific exceptions. Keep try bodies minimal.
try:
state_dict = torch.load(path)
except FileNotFoundError:
raise ValueError(f"Checkpoint not found at {path}") from None
else:
result = convert(state_dict)
# Don't
def make_config(*args):
x, y = args
return dict(**locals())
# Do
def make_config(x, y):
return {"x": x, "y": y}
dataclasses or NamedTuple for configuration objects.Add to all new Python files and shell scripts (not test files). Use the current year.
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
print() — use logger or print_rank_0Any in public APIs, no bare # type: ignore