Linting and Formatting

Style Guides

• Python: Google Python Style Guide
• Shell: Google Shell Style Guide

Formatting and Linting

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

git add -u
pre-commit run
# if it auto-fixed files:
git add -u
pre-commit run

Ruff Rules (from

ruff.toml

)

--fix

--fix

__init__.py

l

O

I

#

• __init__.py

  : F401 and F403 ignored (re-exports expected).

• test_*.py

  ,

  *_test.py

  ,

  tests/*.py

  : D101 and D103 ignored.

Naming Conventions

model_bridge.py

MegatronModelBridge

load_weights_hf_to_megatron

megatron_weights

k

k_99th_percentile

G

G_LOGGER

DEFAULT_HIDDEN_SIZE

• Avoid shadowing variables from an outer scope.
• Initialize all externally visible class members in the constructor.

Import Order

1. __future__

   imports
2. Standard library
3. Third-party (

   megatron.core

   ,

   torch

   ,

   transformers

   , etc.)
4. First-party (

   megatron.bridge.*

   )
5. Local folder imports

I

Type Hints

• Use

  T | None

  instead of

  Optional[T]

• Use

  X | Y

  instead of

  Union[X, Y]

• Use built-in generics (

  list

  ,

  dict

  ,

  tuple

  ) instead of

  typing

  equivalents

def get_module_by_name(
    model: torch.nn.Module,
    name: str,
    default: torch.nn.Module | None = None,
) -> torch.nn.Module | None:
    ...

Mypy

uv run mypy --strict path/to/file.py

• No

  Any

  leaks — use

  object

  for unknown types or a

  TypeVar

  for generics.
• No untyped defs — every function must have parameter and return annotations.
• No implicit

  Optional

  — write

  x: int | None = None

  , never

  x: int = None

  .
• Explicit casts — use

  typing.cast()

  only when inference fails; add a comment.
• Typed dictionaries — prefer

  TypedDict

  over

  dict[str, Any]

  for structured dicts.
• Callable signatures — use

  Callable[[ArgType], ReturnType]

  or

  Protocol

  .
• Ignore sparingly —

  # type: ignore[code]

  must include the error code and justification.

Keyword-Only Arguments for Ambiguous Parameters

*

# 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): ...

Docstrings

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.
    """

Comments

• Commented-out code must have a comment explaining why. Otherwise remove it.
• Comments explain non-obvious intent, trade-offs, or constraints — not what the code does.

Logging

logging.getLogger(__name__)

print_rank_0

warn_rank_0

# Don't
print(f"Loading weights for {model_name}")

# Do
logger = logging.getLogger(__name__)
logger.info("Loading weights for %s", model_name)

Error Handling

try:
    state_dict = torch.load(path)
except FileNotFoundError:
    raise ValueError(f"Checkpoint not found at {path}") from None
else:
    result = convert(state_dict)

Avoid Reflection

# Don't
def make_config(*args):
    x, y = args
    return dict(**locals())

# Do
def make_config(x, y):
    return {"x": x, "y": y}

Configuration and Dataclasses

• Use

  dataclasses

  or

  NamedTuple

  for configuration objects.
• Do not add arbitrary defaults — be explicit about required vs optional fields.

NVIDIA Copyright Header

# 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.

Code Review Checklist

1. Copyright header present on new Python files (not test files)
2. Type hints on public functions and methods
3. Docstrings on public classes and functions (Google style)
4. Specific exceptions in try-except blocks
5. No bare

   print()

   — use

   logger

   or

   print_rank_0

6. No hidden defaults in config function parameters
7. Keyword-only args for ambiguous same-type parameters
8. Double quotes for strings
9. Import order follows the 5-group convention
10. No commented-out code without explanation
11. Mypy clean — no untyped defs, no

    Any

    in public APIs, no bare

    # type: ignore