Megatron-LM Python style conventions

Contract

Enforce the Megatron-LM Python style on every function, method, and class you emit: modern X | None / builtin-generic annotations, Google-style docstrings, specific exception catches, 119-char lines, and the standard naming casing. Apply to all Python code written for this codebase — these are non-default choices the base model gets wrong.

Rules

1. Optionals — X | None, never Optional[X]

• Write an optional value as X | None. NEVER Optional[X] and NEVER Union[X, None].
• The | None must appear at the end of the annotation: dict[str, int] | None, not

None | dict[str, int].

• A parameter with a None default still gets the full annotation: tags: set[str] | None = None.

2. Unions — A | B, never Union[A, B]

• Write a union as A | B | C. NEVER Union[A, B].
• Combine with optional as A | B | None.

3. Generics — builtin lowercase, never the typing aliases

• Use the builtin generics: list[int], dict[str, int], tuple[int, ...], set[str],

type[Foo], frozenset[str].

• NEVER use List, Dict, Tuple, Set, Type, FrozenSet from typing.
• Nest them the same way: list[list[int]], dict[str, list[int]], list[dict[str, int]].
• A homogeneous variable-length tuple is tuple[int, ...]; a fixed-shape tuple names each

slot: tuple[int, int], tuple[str, float, bool].

4. Imports — the deprecated typing names must not appear

• Do NOT write from typing import Optional, Union, List, Dict, Tuple, Set, Type.
• Those symbols must not appear anywhere in the file — not in annotations, not in imports.
• typing is still fine for things with no builtin form (Any, Callable, Protocol,

TypeVar, Iterable); only the six aliases above are banned.

5. Full type hints on every public callable

• Every public function and method carries a type hint on every parameter and on the

return. A function that returns nothing is annotated -> None.

• self and cls are not annotated.

6. Docstrings — Google style on every public class/function/method

• One-line summary first, then a blank line, then the sections.
• Args: section listing each parameter by name.
• Returns: section describing the return value (omit only when the function returns None

and that is obvious; prefer including it).

• Raises: section listing each exception type the body can raise — present if and only

if the body can raise.

• Classes get a class-level docstring describing the class; their public methods get their own.

7. Exceptions — catch specific types, never a bare or catch-all except

• NEVER write a bare except:.
• NEVER write except Exception: as the normal catch-all.
• Catch the specific type(s): except KeyError:, except (ValueError, TypeError):,

except FileNotFoundError:.

8. Formatting & naming

• Line length is 119 characters — not 79, not 88, not 100. Wrap only past 119.
• snake_case for functions and variables, PascalCase for classes, UPPER_SNAKE for

module-level constants.

Worked examples

Rule 1 — optionals

BEFORE (base default):

python
from typing import Optional
def find_checkpoint(run_dir: str, step: int) -> Optional[str]:
    ...

AFTER (conforming):

python
def find_checkpoint(run_dir: str, step: int) -> str | None:
    ...

Rule 2 — unions

BEFORE:

python
from typing import Union
def coerce(x: Union[int, float]) -> Union[int, None]:
    ...

AFTER:

python
def coerce(x: int | float) -> int | None:
    ...

Rule 3 — builtin generics

BEFORE:

python
from typing import List, Dict
def merge_shards(shards: List[int]) -> Dict[str, object]:
    ...

AFTER:

python
def merge_shards(shards: list[int]) -> dict[str, object]:
    ...

Rule 4 — no banned imports

BEFORE:

python
from typing import Optional, List, Dict, Tuple

AFTER (the line is simply gone; annotations use | None, list, dict, tuple):

python
# no typing import of the aliased names at all

Rule 5 + 6 — full hints and Google docstring

BEFORE:

python
def batch_tokenize(texts, max_len):
    """Tokenize texts."""
    ...

AFTER:

python
def batch_tokenize(texts: list[str], max_len: int) -> list[list[int]]:
    """Tokenize a batch of strings into token-id lists.

    Args:
        texts: Input strings to tokenize.
        max_len: Maximum number of tokens per sequence.

    Returns:
        One token-id list per input string.
    """
    ...

Rule 6 — Raises: when the body can raise

BEFORE:

python
def validate_shape(shape: tuple[int, ...]) -> tuple[int, ...]:
    """Validate a shape."""
    if any(d <= 0 for d in shape):
        raise ValueError("dimensions must be positive")
    return shape

AFTER:

python
def validate_shape(shape: tuple[int, ...]) -> tuple[int, ...]:
    """Validate that every dimension of a shape is positive.

    Args:
        shape: The tensor shape to validate.

    Returns:
        The shape unchanged when valid.

    Raises:
        ValueError: If any dimension is non-positive.
    """
    if any(d <= 0 for d in shape):
        raise ValueError("dimensions must be positive")
    return shape

Rule 7 — specific exception catch

BEFORE:

python
def parse_config(path: str):
    try:
        return json.load(open(path))
    except:
        return None

AFTER:

python
def parse_config(path: str) -> dict | None:
    """Read and parse a JSON config file.

    Args:
        path: Path to the config file.

    Returns:
        The parsed config, or None if the file is missing.
    """
    try:
        with open(path) as fh:
            return json.load(fh)
    except FileNotFoundError:
        return None

Edge cases & exceptions

• None default, optional type: annotate as set[str] | None = None. The default

value None does not let you drop the annotation.

• Fixed vs. variable tuple: tuple[int, int] (exactly two ints) is different from

tuple[int, ...] (any number). Pick the one the function actually produces.

• Forward refs / self-types: still use | None inside the quoted form,

e.g. def next(self) -> "Node | None":.

• typing.Any, Callable, Protocol, TypeVar are still imported from typing —

only the six collection/optional aliases are banned.

• Re-raising: except SpecificError: raise is fine; the ban is on catching broadly,

not on re-raising a specific type.

• Truly want every error? Even then, prefer the narrowest superclass that fits

(e.g. except OSError:) over except Exception:. A bare except: is never allowed.

• Private one-liner helpers still take the modern annotations; the docstring rule is

strongest for public API but the typing/exception rules apply everywhere.

Do / Don't

• DON'T from typing import Optional; ALWAYS write X | None inline.
• DON'T annotate Dict[str, int]; ALWAYS write dict[str, int].
• DON'T write except: or except Exception:; ALWAYS name the exact exception type.
• DON'T leave a public function with an untyped parameter; ALWAYS hint every parameter and

the return.

• DON'T write a Raises: section for a function that cannot raise; ALWAYS add one when it can.
• DON'T wrap lines at 79/88/100; ALWAYS allow up to 119.
• DON'T use camelCase methods or lowercase class names; ALWAYS snake_case /

PascalCase / UPPER_SNAKE.

Common mistakes

• Reaching for Optional[...] / Union[...] out of habit — the single most common miss.
• Capitalized generics List[int] / Dict[...] with a from typing import ... line.
• A bare except: or except Exception: pass to "be safe."
• A docstring that is just a one-line summary with no Args: / Returns:.
• Omitting the return annotation (especially -> None).
• Adding a Raises: section out of habit when the function never raises, or omitting it

when the function clearly does.

Quick checklist

1. No Optional[ / Union[ anywhere — X | None, A | B.
2. list/dict/tuple/set/type[...], never the typing aliases; no banned imports.
3. Every public param and return annotated; -> None when nothing is returned.
4. Google docstring: summary, blank line, Args:, Returns:, Raises: (iff it raises).
5. Specific except; never bare or except Exception:.
6. ≤119-char lines; snake_case / PascalCase / UPPER_SNAKE.