From 76fdb9d27f9a3d85bd223ecc7de68cad30efc897 Mon Sep 17 00:00:00 2001 From: aws-arashj <179860940+aws-arashj@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:16:41 +0000 Subject: [PATCH] Publish torchtitan-neuron for Neuron release 2.33 --- .pre-commit-config.yaml | 9 - .../numerics_debugging/activation_tracer.py | 909 ------------- .../numerics_debugging/compare_numerics.py | 1151 ----------------- assets/version.txt | 2 +- pyproject.toml | 19 +- run_train.sh | 32 +- tests/README.md | 8 + tests/integration_tests/__init__.py | 10 +- tests/integration_tests/features.py | 59 +- tests/integration_tests/run_tests.py | 9 +- tests/unit_tests/__init__.py | 49 +- tests/unit_tests/conftest.py | 97 ++ .../unit_tests/test_activation_checkpoint.py | 89 +- tests/unit_tests/test_attn_backend.py | 85 ++ tests/unit_tests/test_checkpoint.py | 113 +- tests/unit_tests/test_compiled_grad_clip.py | 131 ++ .../unit_tests/test_context_parallel_rope.py | 217 ++++ tests/unit_tests/test_loss.py | 21 +- tests/unit_tests/test_loss_dtensor.py | 17 +- tests/unit_tests/test_moe_global_lbl.py | 326 +++++ tests/unit_tests/test_moe_tp_input_grads.py | 76 ++ .../test_moe_tp_only_input_grads_numeric.py | 226 ++++ tests/unit_tests/test_profiler.py | 134 +- tests/unit_tests/test_quantization.py | 6 +- tests/unit_tests/test_rope.py | 65 +- .../test_trainer_profiler_ordering.py | 147 +++ torchtitan/__init__.py | 2 +- torchtitan/components/checkpoint.py | 323 ++++- torchtitan/components/metrics.py | 39 +- torchtitan/components/optimizer.py | 19 +- torchtitan/components/quantization/mx.py | 40 +- .../distributed/activation_checkpoint.py | 48 +- torchtitan/distributed/compile.py | 30 +- torchtitan/distributed/context_parallel.py | 50 +- torchtitan/distributed/neuron_ep/README.md | 15 + torchtitan/distributed/neuron_ep/__init__.py | 20 + torchtitan/distributed/neuron_ep/api.py | 284 ++++ torchtitan/distributed/parallel_dims.py | 44 +- torchtitan/distributed/utils.py | 183 ++- torchtitan/experiments/forge/engine.py | 7 +- torchtitan/experiments/forge/example_train.py | 6 +- .../experiments/graph_trainer/common_utils.py | 63 +- .../graph_trainer/deepseek_v3/parallelize.py | 10 +- .../graph_trainer/llama3/__init__.py | 12 +- .../graph_trainer/llama3/parallelize.py | 22 +- .../graph_trainer/precompile_main.py | 10 +- .../graph_trainer/qwen3/parallelize.py | 10 +- .../experiments/graph_trainer/simple_fsdp.py | 103 +- .../graph_trainer/tests/test_numerics.py | 75 +- .../graph_trainer/tests/test_simple_fsdp.py | 277 +++- .../experiments/graph_trainer/trainer.py | 14 +- torchtitan/experiments/rl/actors/trainer.py | 7 +- torchtitan/experiments/torchft/trainer.py | 7 +- .../transformers_modeling_backend/model.py | 23 +- torchtitan/models/common/attention.py | 34 +- torchtitan/models/common/config_utils.py | 55 +- torchtitan/models/common/decoder.py | 10 +- torchtitan/models/common/moe.py | 120 +- torchtitan/models/common/moe_sharding.py | 9 +- torchtitan/models/common/rope.py | 120 +- torchtitan/models/common/token_dispatcher.py | 347 ++++- torchtitan/models/deepseek_v3/sharding.py | 8 +- torchtitan/models/flux/trainer.py | 6 +- torchtitan/models/gpt_oss/__init__.py | 29 +- torchtitan/models/gpt_oss/model.py | 133 +- torchtitan/models/gpt_oss/sharding.py | 15 +- torchtitan/models/qwen3/__init__.py | 1 - torchtitan/models/qwen3/config_registry.py | 18 +- torchtitan/models/qwen3/model.py | 5 +- torchtitan/models/qwen3/sharding.py | 13 +- torchtitan/protocols/checkpoint_load.py | 45 + torchtitan/tools/profiler.py | 164 ++- torchtitan/tools/utils.py | 29 +- torchtitan/trainer.py | 85 +- torchtitan_neuron/README.md | 66 + torchtitan_neuron/__init__.py | 11 + torchtitan_neuron/compile_backend.py | 40 + torchtitan_neuron/components/__init__.py | 4 + .../components/quantization/__init__.py | 33 + .../components/quantization/mx_prequant.py | 544 ++++++++ torchtitan_neuron/docs/dense-tutorials.md | 199 +++ torchtitan_neuron/docs/moe-tutorials.md | 58 + .../docs/neuron_attention_backends.md | 129 ++ torchtitan_neuron/docs/neuron_mxfp8.md | 36 + .../docs/qwen3_235b_full_ep128.md | 147 +++ torchtitan_neuron/docs/qwen3_inference.md | 69 + torchtitan_neuron/docs/qwen3_nki_moe.md | 77 ++ torchtitan_neuron/experiments/__init__.py | 6 + .../experiments/graph_trainer/__init__.py | 6 + .../experiments/graph_trainer/configs.py | 43 + .../graph_trainer/deepseek_v3/__init__.py | 6 + .../deepseek_v3/config_registry.py | 34 + .../graph_trainer/gpt_oss/__init__.py | 38 + .../graph_trainer/gpt_oss/config_registry.py | 26 + .../graph_trainer/gpt_oss/model.py | 24 + .../graph_trainer/gpt_oss/parallelize.py | 77 ++ .../graph_trainer/llama3/__init__.py | 6 + .../graph_trainer/llama3/config_registry.py | 70 + .../graph_trainer/qwen3/__init__.py | 41 + .../graph_trainer/qwen3/config_registry.py | 375 ++++++ torchtitan_neuron/models/__init__.py | 4 + .../models/deepseek_v3/__init__.py | 6 + .../models/deepseek_v3/config_registry.py | 71 + .../models/deepseek_v3/nki_moe.py | 658 ++++++++++ torchtitan_neuron/models/gpt_oss/__init__.py | 4 + .../models/gpt_oss/config_registry.py | 63 + .../models/gpt_oss/nki_attention.py | 339 +++++ torchtitan_neuron/models/gpt_oss/nki_moe.py | 760 +++++++++++ torchtitan_neuron/models/llama3/__init__.py | 6 + .../models/llama3/config_registry.py | 97 ++ torchtitan_neuron/models/qwen3/__init__.py | 99 ++ .../models/qwen3/_mxfp8_flavors.py | 127 ++ .../models/qwen3/config_registry.py | 670 ++++++++++ torchtitan_neuron/models/qwen3/nki_moe.py | 795 ++++++++++++ torchtitan_neuron/tests/README.md | 23 + torchtitan_neuron/tests/__init__.py | 38 + torchtitan_neuron/tests/conftest.py | 12 + torchtitan_neuron/tests/models/__init__.py | 4 + .../tests/models/gpt_oss/__init__.py | 4 + .../test_graph_trainer_config_registry.py | 63 + .../gpt_oss/test_trainer_config_registry.py | 83 ++ .../tests/models/llama3/__init__.py | 4 + .../test_graph_trainer_config_registry.py | 69 + .../llama3/test_trainer_config_registry.py | 54 + .../tests/models/qwen3/__init__.py | 4 + .../tests/models/qwen3/test_global_lbl.py | 103 ++ .../test_graph_trainer_config_registry.py | 198 +++ .../tests/models/qwen3/test_nki_moe.py | 483 +++++++ .../qwen3/test_trainer_config_registry.py | 237 ++++ .../tests/moe_block_microbench.py | 463 +++++++ .../tests/moe_dispatcher_microbench.py | 567 ++++++++ .../tests/nki_blockwise_moe_cte.py | 38 + .../tests/test_attention_masks.py | 98 ++ torchtitan_neuron/tests/test_capability.py | 80 ++ .../tests/test_compile_backend.py | 47 + .../tests/test_dense_config_registry.py | 143 ++ torchtitan_neuron/tests/test_ep_dispatcher.py | 929 +++++++++++++ torchtitan_neuron/tests/test_ep_kernels.py | 223 ++++ .../tests/test_flex_attention.py | 139 ++ torchtitan_neuron/tests/test_metrics_mfu.py | 49 + .../tests/test_mxfp8_numerics.py | 238 ++++ 141 files changed, 14853 insertions(+), 2400 deletions(-) delete mode 100644 agent_tooling/numerics_debugging/activation_tracer.py delete mode 100644 agent_tooling/numerics_debugging/compare_numerics.py create mode 100644 tests/unit_tests/conftest.py create mode 100644 tests/unit_tests/test_attn_backend.py create mode 100644 tests/unit_tests/test_compiled_grad_clip.py create mode 100644 tests/unit_tests/test_context_parallel_rope.py create mode 100644 tests/unit_tests/test_moe_global_lbl.py create mode 100644 tests/unit_tests/test_moe_tp_input_grads.py create mode 100644 tests/unit_tests/test_moe_tp_only_input_grads_numeric.py create mode 100644 tests/unit_tests/test_trainer_profiler_ordering.py create mode 100644 torchtitan/distributed/neuron_ep/README.md create mode 100644 torchtitan/distributed/neuron_ep/__init__.py create mode 100644 torchtitan/distributed/neuron_ep/api.py create mode 100644 torchtitan/protocols/checkpoint_load.py create mode 100644 torchtitan_neuron/README.md create mode 100644 torchtitan_neuron/__init__.py create mode 100644 torchtitan_neuron/compile_backend.py create mode 100644 torchtitan_neuron/components/__init__.py create mode 100644 torchtitan_neuron/components/quantization/__init__.py create mode 100644 torchtitan_neuron/components/quantization/mx_prequant.py create mode 100644 torchtitan_neuron/docs/dense-tutorials.md create mode 100644 torchtitan_neuron/docs/moe-tutorials.md create mode 100644 torchtitan_neuron/docs/neuron_attention_backends.md create mode 100644 torchtitan_neuron/docs/neuron_mxfp8.md create mode 100644 torchtitan_neuron/docs/qwen3_235b_full_ep128.md create mode 100644 torchtitan_neuron/docs/qwen3_inference.md create mode 100644 torchtitan_neuron/docs/qwen3_nki_moe.md create mode 100644 torchtitan_neuron/experiments/__init__.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/__init__.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/configs.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/deepseek_v3/__init__.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/deepseek_v3/config_registry.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/gpt_oss/__init__.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/gpt_oss/config_registry.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/gpt_oss/model.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/gpt_oss/parallelize.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/llama3/__init__.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/llama3/config_registry.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/qwen3/__init__.py create mode 100644 torchtitan_neuron/experiments/graph_trainer/qwen3/config_registry.py create mode 100644 torchtitan_neuron/models/__init__.py create mode 100644 torchtitan_neuron/models/deepseek_v3/__init__.py create mode 100644 torchtitan_neuron/models/deepseek_v3/config_registry.py create mode 100644 torchtitan_neuron/models/deepseek_v3/nki_moe.py create mode 100644 torchtitan_neuron/models/gpt_oss/__init__.py create mode 100644 torchtitan_neuron/models/gpt_oss/config_registry.py create mode 100644 torchtitan_neuron/models/gpt_oss/nki_attention.py create mode 100644 torchtitan_neuron/models/gpt_oss/nki_moe.py create mode 100644 torchtitan_neuron/models/llama3/__init__.py create mode 100644 torchtitan_neuron/models/llama3/config_registry.py create mode 100644 torchtitan_neuron/models/qwen3/__init__.py create mode 100644 torchtitan_neuron/models/qwen3/_mxfp8_flavors.py create mode 100644 torchtitan_neuron/models/qwen3/config_registry.py create mode 100644 torchtitan_neuron/models/qwen3/nki_moe.py create mode 100644 torchtitan_neuron/tests/README.md create mode 100644 torchtitan_neuron/tests/__init__.py create mode 100644 torchtitan_neuron/tests/conftest.py create mode 100644 torchtitan_neuron/tests/models/__init__.py create mode 100644 torchtitan_neuron/tests/models/gpt_oss/__init__.py create mode 100644 torchtitan_neuron/tests/models/gpt_oss/test_graph_trainer_config_registry.py create mode 100644 torchtitan_neuron/tests/models/gpt_oss/test_trainer_config_registry.py create mode 100644 torchtitan_neuron/tests/models/llama3/__init__.py create mode 100644 torchtitan_neuron/tests/models/llama3/test_graph_trainer_config_registry.py create mode 100644 torchtitan_neuron/tests/models/llama3/test_trainer_config_registry.py create mode 100644 torchtitan_neuron/tests/models/qwen3/__init__.py create mode 100644 torchtitan_neuron/tests/models/qwen3/test_global_lbl.py create mode 100644 torchtitan_neuron/tests/models/qwen3/test_graph_trainer_config_registry.py create mode 100644 torchtitan_neuron/tests/models/qwen3/test_nki_moe.py create mode 100644 torchtitan_neuron/tests/models/qwen3/test_trainer_config_registry.py create mode 100644 torchtitan_neuron/tests/moe_block_microbench.py create mode 100644 torchtitan_neuron/tests/moe_dispatcher_microbench.py create mode 100644 torchtitan_neuron/tests/nki_blockwise_moe_cte.py create mode 100644 torchtitan_neuron/tests/test_attention_masks.py create mode 100644 torchtitan_neuron/tests/test_capability.py create mode 100644 torchtitan_neuron/tests/test_compile_backend.py create mode 100644 torchtitan_neuron/tests/test_dense_config_registry.py create mode 100644 torchtitan_neuron/tests/test_ep_dispatcher.py create mode 100644 torchtitan_neuron/tests/test_ep_kernels.py create mode 100644 torchtitan_neuron/tests/test_flex_attention.py create mode 100644 torchtitan_neuron/tests/test_metrics_mfu.py create mode 100644 torchtitan_neuron/tests/test_mxfp8_numerics.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 180a34e6008..bb6153b4369 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,15 +17,6 @@ repos: - id: end-of-file-fixer exclude: '^(.*\.svg)$' -- repo: https://github.com/Lucas-C/pre-commit-hooks - rev: v1.5.4 - hooks: - - id: insert-license - files: \.py$ - args: - - --license-filepath - - assets/license_header.txt - - repo: https://github.com/pycqa/flake8 rev: 34cbf8ef3950f43d09b85e2e45c15ae5717dc37b hooks: diff --git a/agent_tooling/numerics_debugging/activation_tracer.py b/agent_tooling/numerics_debugging/activation_tracer.py deleted file mode 100644 index 8ae7eb1c4db..00000000000 --- a/agent_tooling/numerics_debugging/activation_tracer.py +++ /dev/null @@ -1,909 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -""" -Numerics activation tracer built on top of ``torch.utils._debug_mode.DebugMode``. - -Captures per-op activation snapshots during training without modifying -model source code. Each capture carries the originating module FQN, -source location, phase (forward/backward), and norm-hashes of input + -output tensors (for detecting in-place mutations and bit-exact -divergences). Works with DTensor, torch.compile, FSDP, and activation -checkpointing. - -For each captured op, we track: - - **Op name**: the ATen operation (e.g. mm, bmm, silu, add) - - **Module FQN**: which module the op belongs to - (e.g. ``layers.0.attention.qkv_linear.wq``) - - **Source location**: stack trace back to the model code that - triggered the op (e.g. ``common/attention.py:534``) - - **Phase**: whether the op ran during forward or backward - - **Output hash / Input hashes**: DebugMode's ``norm_hash_fn`` - applied to result and pre-op inputs. - -What gets captured: - By default, captures outputs of every op that isn't in ``_EXCLUDED_OPS`` - (denylist of infrastructure ops: view, reshape, transpose, clone, copy_, - detach, collective comms, etc.) and that produces a float tensor with at - least ``min_numel`` elements (default 1000). This errs on the side of - capturing too much rather than silently missing newly introduced compute - ops (e.g. ``_scaled_dot_product_flash_attention``, ``_scaled_mm``). - - Each captured tensor is keyed by ``module_fqn/op_N_opname`` (e.g. - ``layers.0.attention.qkv_linear.wq/op_0_mm``) and tagged with its - phase (forward or backward). - -Filtering / customization: - Pass ``op_filter`` to ``DebugModeTracer`` to restrict capture to specific - ops (matched by substring against the op name). For example, to only - capture matrix multiplies and softmax:: - - DebugModeTracer(model, op_filter={"mm", "bmm", "softmax"}) - - Tweak ``min_numel`` to filter small tensors, and edit ``_EXCLUDED_OPS`` - to drop additional infrastructure ops you never want to see. -""" - -import os -import re -from contextvars import ContextVar -from dataclasses import dataclass, field - -import torch -import torch.utils._pytree as pytree -from torch import nn - - -@dataclass -class _StackFrame: - filename: str - lineno: int - name: str - line: str | None = None - - def short_str(self) -> str: - parts = self.filename.replace("\\", "/").split("/") - short_path = "/".join(parts[-2:]) if len(parts) >= 2 else self.filename - return f"{short_path}:{self.lineno}" - - -@dataclass -class CapturedActivation: - """A captured activation: pre-computed stats + provenance. - - Stats (Shape, Dtype, L2 norm, Min/Max/Mean) are computed inline at - capture time so we never hold a tensor clone after the op completes. - L1 is represented by DebugMode's output hash. For in-place ops where - the dispatch returns ``None`` the dict is empty — those entries rely - on ``input_hashes`` / ``output_hash`` for comparison. - """ - - # Pre-computed stat strings keyed by name (``Shape``, ``L2 norm``, - # etc.). Empty for in-place op captures (no result to summarize). - stats: dict[str, str] = field(default_factory=dict) - stack_frames: list[_StackFrame] = field(default_factory=list) - phase: str = "forward" - # DebugMode norm-hash (basically L1) of the op result. Comma- - # separated when the result is a tuple/list of tensors. Empty - # string means no hash recorded. - output_hash: str = "" - # Comma-separated norm-hashes of input tensors captured *before* - # the op ran. For in-place ops (e.g. ``_fused_adam_``) comparing - # an op's ``input_hashes`` against the previous op's - # ``input_hashes`` for the same arg surfaces the mutation, even - # though ``output_hash`` is empty. - input_hashes: str = "" - # Producing capture key for each input tensor, aligned positionally - # with ``input_hashes``. Resolved from DebugMode's tensor ID - # tracker: if an input tensor was the output of an earlier captured - # op, that op's key is stored; otherwise empty. Semicolon- - # separated (commas are reserved for the hash list). - input_producers: str = "" - # The op's aten name (e.g. ``_fused_adam_``). Persisted so the - # comparison tool can recognize in-place ops without re-parsing - # the key. - op_name: str = "" - - -def _compute_stats(tensor: torch.Tensor) -> dict[str, str]: - """Reduce ``tensor`` to the stat strings written to the log. - - Done inline at capture time so we never have to clone or hold - onto the tensor itself. L2 norm and mean are reduced in - **float64** so they're deterministic across reduction-order drift - (eager vs traced reorder sums differently — float32 reductions - diverge in the last digit; float64 has ~9 more bits of precision - and absorbs it). L1 norm is intentionally omitted — it's - identical to ``output_hash`` (which DebugMode computes via - ``norm_hash_fn`` in float64), so dumping it again would be - redundant. Min/Max are exact, so we keep the original dtype. - """ - t64 = tensor.to(dtype=torch.float64) - return { - "Shape": f"{tensor.shape}, Dtype: {tensor.dtype}", - "L2 norm": f"{t64.norm(p=2).item():.6e}", - "Min": f"{tensor.float().min().item():.6e}", - "Max": f"{tensor.float().max().item():.6e}", - "Mean": f"{t64.mean().item():.6e}", - } - - -def _stringify_hash_tree(hash_obj) -> str: - """Flatten a DebugMode hash pytree into a comma-separated string. - - ``log_tensor_hashes`` stores hashes as a pytree mirroring the - original args / kwargs / result structure, with floats where - tensors were and ``None`` elsewhere. We flatten + drop Nones so - the log line is human-readable. - - Returns an empty string when nothing hashable was present. - """ - if hash_obj is None: - return "" - leaves = pytree.tree_leaves(hash_obj) - vals = [v for v in leaves if v is not None] - if not vals: - return "" - return ", ".join(f"{v:.6e}" if isinstance(v, float) else str(v) for v in vals) - - -# Ops to always exclude — infrastructure, not meaningful for numerics parity. -# These are creation, view/reshape, indexing, and collective comm ops that -# either don't perform compute or just rewrap the underlying storage. -_EXCLUDED_OPS = frozenset( - { - # Creation - "empty", - "empty_like", - "zeros", - "zeros_like", - "ones", - "ones_like", - "full", - "full_like", - "lift_fresh", - "lift_fresh_copy", - # Copies / dtype conversion - "copy_", - "_to_copy", - "to", - "clone", - "contiguous", - "_pin_memory", - # View / reshape - "view", - "_unsafe_view", - "view_as", - "view_as_real", - "view_as_complex", - "reshape", - "_reshape_alias", - "flatten", - "unflatten", - "alias", - # Permutation - "permute", - "transpose", - "t", - "unsqueeze", - "squeeze", - "expand", - "expand_as", - "repeat", - # Concatenation / splitting - "cat", - "stack", - "split", - "chunk", - # Indexing - "narrow", - "select", - "index_select", - "gather", - "scatter", - "slice", - "as_strided", - # Autograd plumbing - "detach", - # Collective comm (DTensor / FSDP). - # - # Async collective ops dispatch and return *before* the - # destination buffer is filled. - # "redistribute", - # "all_gather", - # "all_gather_into_tensor", - # "all_gather_into_tensor_out", - "_pre_bucket_all_gather", - # "all_reduce", - # "reduce_scatter", - # "reduce_scatter_tensor", - "_pre_bucket_reduce_scatter", - } -) - - -def _is_filtered_frame(path: str) -> bool: - """Filter frames from torch internals and this module. - - Used by ``_parse_stack_trace`` when surfacing source locations - from ``node.meta["stack_trace"]`` and from DebugMode's - ``op.fwd_stack_trace`` / ``op.stack_trace`` strings. - """ - p = path.replace("\\", "/") - if "/torch/nn/" in p: - return True - if "/torch/autograd/" in p: - return True - # Dispatcher / compile / runtime plumbing — these frames show up - # in DebugMode's live traceback under FSDP and activation - # checkpoint, drowning out the user-code frame. - if "/torch/_ops.py" in p: - return True - if "/torch/_compile.py" in p: - return True - if "/torch/_dynamo/" in p: - return True - if "/torch/utils/_debug_mode/" in p: - return True - if "/torch/utils/checkpoint.py" in p: - return True - if "activation_tracer.py" in p: - return True - # FQNInterpreter wraps every traced node's execution, so its - # super().run_node line shows up in the live traceback when a - # node has no node.meta["stack_trace"]. It's never the - # location the user wants. - if "graph_trainer/debug_utils.py" in p: - return True - return False - - -def _get_op_name(func) -> str: - if hasattr(func, "_schema"): - return func._schema.name.split("::")[-1] - return str(func).split(".")[-1].replace(">", "") - - -def _local_tensor_if_dtensor(tensor: torch.Tensor) -> torch.Tensor: - try: - from torch.distributed.tensor import DTensor - - return tensor._local_tensor if isinstance(tensor, DTensor) else tensor - except ImportError: - return tensor - - -def _is_captureable_tensor(tensor: torch.Tensor, min_numel: int) -> bool: - # Skip FakeTensors produced during make_fx tracing. - if isinstance(tensor, torch._subclasses.FakeTensor): - return False - - tensor = _local_tensor_if_dtensor(tensor) - if tensor.numel() < min_numel: - return False - return tensor.dtype in {torch.float32, torch.float16, torch.bfloat16} - - -def _captureable_tensor_leaves(result, min_numel: int) -> list[torch.Tensor]: - """Return result tensor leaves that are meaningful for numerics capture.""" - tensors: list[torch.Tensor] = [] - for leaf in pytree.tree_leaves(result): - if isinstance(leaf, torch.Tensor) and _is_captureable_tensor(leaf, min_numel): - tensors.append(_local_tensor_if_dtensor(leaf)) - return tensors - - -def _should_capture(func, result, min_numel: int, op_filter: set[str] | None) -> bool: - """Filter to only capture "interesting" tensor outputs. - - Filters by: - - Tensor type (must be a real tensor, not FakeTensor) - - Size (must have at least min_numel elements) - - Dtype (must be float32, float16, or bfloat16) - - Always excludes infrastructure ops (empty, copy_, view, etc.) - - Optionally by op name (if op_filter is provided) - - For DTensors, the local tensor is used for size/dtype checks. - - Args: - func: The dispatch function. - result: The operation result. - min_numel: Minimum number of elements to capture. - op_filter: Optional set of op name substrings to match. - If None, captures all ops that pass other filters. - - Returns: - True if the result should be captured. - """ - op_name = _get_op_name(func) - if op_name in _EXCLUDED_OPS: - return False - if op_filter is None: - return bool(_captureable_tensor_leaves(result, min_numel)) - return any(f in op_name for f in op_filter) and bool( - _captureable_tensor_leaves(result, min_numel) - ) - - -# The current module FQN for the op being dispatched. -# -# Set by ``FQNInterpreter.run_node`` during traced graph replay (from -# ``node.meta["custom"]["module_fqn"]``). Read by -# ``DebugModeTracer.record_hook`` to build the capture key (e.g. -# ``"layers.0.attention.wq/op_0_mm"``). -# -# During eager mode this stays ``None``; the FQN is recovered from -# DebugMode's ``ModTracker`` (for forward ops) or from -# ``_grad_fn_to_module`` (for backward ops). -_current_module_name: ContextVar[str | None] = ContextVar( - "_current_module_name", default=None -) - -# Override for stack frames attached to captured activations. -# -# Set by ``FQNInterpreter`` (from ``node.meta["stack_trace"]``) during -# traced replay. When set, ``record_hook`` uses these instead of -# DebugMode's live or autograd-forward traceback strings. -_current_stack_frames: ContextVar[list[_StackFrame] | None] = ContextVar( - "_current_stack_frames", default=None -) - -# Phase override for the current op ("forward" or "backward"). -# -# Set by ``FQNInterpreter`` from ``node.meta["autograd_backward"]`` -# during traced replay. When set, ``record_hook`` uses this directly -# instead of inferring the phase from ``torch._C._current_autograd_node()``. -_current_phase_override: ContextVar[str | None] = ContextVar( - "_current_phase_override", default=None -) - -# Global flag indicating that ActivationCaptureProfiler has armed capture. -# -# Checked by GraphTrainer._maybe_get_fqn_interpreter() to decide whether to -# route traced graph replay through FQNInterpreter (which populates -# _current_module_name from node metadata) instead of direct gm(*inputs). -# Module-level bool rather than a callback — each rank is a separate -# process, so no thread-safety concern. -_numerics_capture_active = False - - -def is_numerics_capture_active() -> bool: - return _numerics_capture_active - - -def set_numerics_capture_active(active: bool) -> None: - global _numerics_capture_active - _numerics_capture_active = active - - -_STACK_TRACE_RE = re.compile(r'File "([^"]+)", line (\d+), in (\S+)\n\s*(.*?)(?=\n|$)') - - -def _parse_stack_trace(trace_str: str) -> list[_StackFrame]: - """Parse File/line/func entries from node.meta["stack_trace"], - filtering out torch.nn.modules frames.""" - frames = [] - for m in _STACK_TRACE_RE.finditer(trace_str): - if not _is_filtered_frame(m.group(1)): - frames.append( - _StackFrame( - filename=m.group(1), - lineno=int(m.group(2)), - name=m.group(3), - line=m.group(4).strip() or None, - ) - ) - return frames - - -def _clean_fqn(fqn: str) -> str: - """Strip activation-checkpoint wrapper names from FQNs.""" - return fqn.replace("._checkpoint_wrapped_module", "") - - -# Maps autograd Node (grad_fn) to module FQN. Built during forward -# by the global module forward hooks installed in -# ``DebugModeTracer.__enter__`` (via ``_record_grad_fn``), read during -# backward by ``record_hook`` to give backward ops their originating -# module's FQN. -_grad_fn_to_module: dict[object, str] = {} - - -def _collect_grad_fns(args, kwargs) -> set: - """Collect grad_fn objects from input tensors.""" - grad_fns = set() - for a in pytree.tree_leaves((args, kwargs)): - if isinstance(a, torch.Tensor) and a.grad_fn is not None: - grad_fns.add(a.grad_fn) - return grad_fns - - -def _record_grad_fn(result, module_name: str, input_grad_fns: set) -> None: - """Map output grad_fn to module FQN for backward op attribution. - - Walks the autograd graph from the output grad_fn (DFS via an - explicit stack), mapping internal nodes to the module. The walk - stops at input grad_fns (sibling/parent module boundary). It does - not overwrite nodes already mapped to child modules, but it still - walks through them so parent modules can claim unmapped glue ops - between children. Traversal order doesn't matter for correctness - because the "claim only if unclaimed" check is order-independent. - """ - for t in pytree.tree_leaves(result): - if not isinstance(t, torch.Tensor) or t.grad_fn is None: - continue - stack = [t.grad_fn] - visited = set() - while stack: - fn = stack.pop() - if fn in visited or fn in input_grad_fns: - continue - visited.add(fn) - if fn not in _grad_fn_to_module: - _grad_fn_to_module[fn] = module_name - # Walk through child-module nodes (already mapped) to - # reach unmapped nodes behind them — those belong to - # this (parent) module. - for parent, _ in fn.next_functions: - if parent is not None and parent not in visited: - stack.append(parent) - - -def dump_captures_to_file( - captures: dict[str, CapturedActivation], - filepath: str, - *, - skipped_excluded_ops: set[str] | None = None, -) -> None: - """Write per-op activation statistics to a human-readable log file. - - Args: - captures: Captured activations to dump. - filepath: Output path. - skipped_excluded_ops: Optional set of op names that dispatched - during the captured step but were dropped because they - live in ``_EXCLUDED_OPS``. When provided, written into the - log header as a comma-separated list so the comparison - tool can surface them in the HTML. - """ - os.makedirs(os.path.dirname(filepath), exist_ok=True) - with open(filepath, "w") as f: - f.write(f"Total captured activations: {len(captures)}\n") - if skipped_excluded_ops is not None: - ops_str = ", ".join(sorted(skipped_excluded_ops)) or "(none)" - f.write(f"Excluded ops dispatched: {ops_str}\n") - f.write("=" * 80 + "\n\n") - for key, cap in captures.items(): - f.write(f"[{_clean_fqn(key)}]\n") - if cap.stats: - # Stats are pre-computed at capture time; the dump is - # just a write-through. Order matters for grep-ability, - # so use a fixed sequence rather than dict iteration. - for stat in ("Shape", "L1 norm", "L2 norm", "Min", "Max", "Mean"): - val = cap.stats.get(stat) - if val is not None: - f.write(f" {stat}: {val}\n") - if cap.output_hash: - f.write(f" Output hash: {cap.output_hash}\n") - if cap.input_hashes: - f.write(f" Input hashes: {cap.input_hashes}\n") - if cap.input_producers: - f.write(f" Input producers: {cap.input_producers}\n") - if cap.stack_frames: - f.write(f" Location: {cap.stack_frames[-1].short_str()}\n") - if cap.phase != "forward": - f.write(f" Phase: {cap.phase}\n") - f.write("\n") - - -class DebugModeTracer: - """Captures per-op activations via :class:`torch.utils._debug_mode.DebugMode`. - - Drives the capture by entering a DebugMode with - ``record_nn_module=True`` (for module FQNs via ``ModTracker``), - ``record_stack_trace=True`` (for per-op source locations + autograd - forward tracebacks), ``record_ids=True`` (graph-style tensor IDs - used for input-producer attribution), and - ``log_tensor_hashes(hash_inputs=True)`` (norm hashes of inputs and - outputs, including for in-place ops where the result is ``None``). - - Per-op metadata is stamped by a custom dispatch hook into - ``op.record``; after ``__exit__`` the captured operator stream is - walked once to materialize a ``dict[str, CapturedActivation]`` - keyed by ``module_fqn/op_N_opname`` and consumable by - :mod:`agent_tooling.numerics_debugging.compare_numerics`. - - Args: - model: The model whose activations to capture. Used to install - global ``nn.Module`` forward hooks that populate - :data:`_grad_fn_to_module` so backward ops can recover the - originating module's FQN. - min_numel: Minimum tensor size to capture (default 1000). - op_filter: Optional set of op-name substrings (allowlist). None - captures every op not in ``_EXCLUDED_OPS``. - """ - - def __init__( - self, - model: nn.Module, - *, - min_numel: int = 1000, - op_filter: set[str] | None = None, - ): - self.model = model - self.min_numel = min_numel - self.op_filter = op_filter - self._captures: dict[str, CapturedActivation] = {} - self.skipped_excluded_ops: set[str] = set() - self._debug_mode = None - self._hook_ctx = None - self._hash_ctx = None - self._materialized = False - # Global nn.Module forward hook handles, removed on __exit__. - self._pre_fw_handle = None - self._post_fw_handle = None - - def __enter__(self) -> dict[str, CapturedActivation]: - from torch.utils._debug_mode import DebugMode - - self._debug_mode = DebugMode( - record_realtensor=True, - record_faketensor=False, - record_nn_module=True, - record_stack_trace=True, - # record_ids populates DebugMode's _tensor_memo with a stable - # ``$N`` graph ID per tensor. We share that memo to map - # input tensors back to the captured op that produced them - # (used for the producer-link column in the diff HTML). - record_ids=True, - ) - # log_tensor_hashes stamps op.log["hash"] (norm hash of the - # result, basically L1) and op.log["input_hash"] (norm hash of - # each input tensor *before* the op runs). Input hashes are - # what catches in-place mutations — ``_fused_adam_`` returns - # None so its output is unhashable, but the per-param input - # hashes are still recorded. - self._hash_ctx = DebugMode.log_tensor_hashes(hash_fn="norm", hash_inputs=True) - - # Install global forward hooks so backward ops can recover - # their owning module's FQN via the autograd graph. DebugMode's - # ModTracker only fires on Python forward calls, so during - # backward (driven by the C++ autograd engine) its module stack - # is empty. These hooks fire on every nn.Module.__call__: - # walk the autograd graph from each module's output during the - # forward post-hook and stash grad_fn -> fqn in _grad_fn_to_module. - # _maybe_record_module_fqn in record_hook reads it during backward. - _grad_fn_to_module.clear() - module_to_fqn = {id(m): name for name, m in self.model.named_modules()} - pending_input_grad_fns: dict[int, set] = {} - - def pre_fw_hook(module, args): - mid = id(module) - if mid not in module_to_fqn: - return - pending_input_grad_fns[mid] = _collect_grad_fns(args, {}) - - def post_fw_hook(module, args, output): - mid = id(module) - fqn = module_to_fqn.get(mid) - if fqn is None: - return - input_grad_fns = pending_input_grad_fns.pop(mid, set()) - _record_grad_fn(output, fqn, input_grad_fns) - - self._pre_fw_handle = nn.modules.module.register_module_forward_pre_hook( - pre_fw_hook - ) - self._post_fw_handle = nn.modules.module.register_module_forward_hook( - post_fw_hook - ) - - # Computes output stats only for ops that pass the filter — - # cloning every dispatch output (including AC recompute) blows - # out CUDA memory. We also stash module FQN, stack frames, and - # phase per op so the materialize step doesn't need to walk - # DebugMode annotate entries. - skipped_excluded_ops = self.skipped_excluded_ops - min_numel = self.min_numel - op_filter = self.op_filter - debug_mode = self._debug_mode - - def record_hook(func, types, args, kwargs, result): - out = {} - op_name = _get_op_name(func) - if op_name in _EXCLUDED_OPS: - skipped_excluded_ops.add(op_name) - return out - - # Compute input/output tensor IDs using DebugMode's shared - # tensor memo — lets _materialize_captures map an input - # back to its producing op. Done unconditionally for all - # non-excluded ops (including in-place ops we won't clone) - # so the producer map is complete. - tracker = debug_mode._tensor_memo - - def _id_leaves(obj): - tree = pytree.tree_map( - lambda x: tracker._id(x) if isinstance(x, torch.Tensor) else None, - obj, - ) - return [v for v in pytree.tree_leaves(tree) if v is not None] - - out["input_ids"] = _id_leaves((args, kwargs)) - out["output_ids"] = _id_leaves(result) - - if not _should_capture(func, result, min_numel, op_filter): - # Still want phase + module FQN + stack frames for the - # in-place op path in _materialize_captures. Fall - # through to the metadata block below. - pass - else: - # Compute stats for the first captureable result tensor. - # Multi-output ops still get complete DebugMode hash - # strings for all tensor leaves; the scalar stat columns - # summarize the primary tensor leaf so the existing log - # format stays stable. - primary_tensor = _captureable_tensor_leaves(result, min_numel)[0] - # Compute stats inline — no clone, no held reference. - out["stats"] = _compute_stats(primary_tensor) - - # Module FQN priority: - # 1. _current_module_name ContextVar — set by FQNInterpreter - # during traced graph replay (gm bypasses module forwards - # so ModTracker is silent). - # 2. DebugMode's ModTracker stack — populated by Python - # module forward hooks in eager mode. ModTracker roots - # FQNs at the model class (``FSDPLlama3Model.layers.0...``) - # while named_modules() does not — strip the root segment - # so all three sources agree on the same naming. - # 3. _grad_fn_to_module — populated by our forward hook - # (above); recovers FQN for backward ops, where the C++ - # autograd engine bypasses module forwards so ModTracker - # is silent. - fqn = _current_module_name.get() - if fqn is None and debug_mode.current_nn_module_stack: - fqn = debug_mode.current_nn_module_stack[-1] - root = debug_mode.current_nn_module_stack[0] - if fqn == root: - fqn = None - elif fqn.startswith(root + "."): - fqn = fqn[len(root) + 1 :] - if fqn is None: - autograd_node = torch._C._current_autograd_node() - if autograd_node is not None: - fqn = _grad_fn_to_module.get(autograd_node) - if fqn: - out["module_fqn"] = _clean_fqn(fqn) - - # Stack frames: ContextVar (from FQNInterpreter, already - # parsed) takes priority over the live traceback DebugMode - # captured at dispatch time, which is unhelpful under FSDP / - # graph replay. - frames = _current_stack_frames.get() - if frames is not None: - out["stack_frames"] = frames - - phase = _current_phase_override.get() - if phase is None: - autograd_node = torch._C._current_autograd_node() - phase = "backward" if autograd_node is not None else "forward" - out["phase"] = phase - return out - - self._hook_ctx = DebugMode.dispatch_hooks(record_hook=record_hook) - self._debug_mode.__enter__() - self._hash_ctx.__enter__() - self._hook_ctx.__enter__() - return self._captures - - def __exit__(self, *args): - try: - if self._hook_ctx is not None: - hook_ctx = self._hook_ctx - self._hook_ctx = None - hook_ctx.__exit__(*args) - if self._hash_ctx is not None: - hash_ctx = self._hash_ctx - self._hash_ctx = None - hash_ctx.__exit__(*args) - if self._debug_mode is not None: - self._debug_mode.__exit__(*args) - if self._debug_mode is not None and not self._materialized: - self._materialize_captures() - self._materialized = True - finally: - if self._pre_fw_handle is not None: - self._pre_fw_handle.remove() - self._pre_fw_handle = None - if self._post_fw_handle is not None: - self._post_fw_handle.remove() - self._post_fw_handle = None - _grad_fn_to_module.clear() - # Mark the DebugMode context as closed even if materializing - # captures raises, so repeated teardown cannot double-exit it. - self._debug_mode = None - - def _materialize_captures(self) -> None: - """Build ``module_fqn/op_N_opname`` keyed captures from the - per-op records stamped by ``record_hook``. - - Most metadata (FQN, stack, phase, output clone) is stamped at - hook time, so this pass just assigns per-module op indices and - threads through the tensor-ID producer map used for input-link - attribution in the diff HTML. - """ - try: - from torch.utils._debug_mode import _OpCall - except ImportError: - from torch.utils._debug_mode._calls import _OpCall - - op_counters: dict[str, int] = {} - # Map from tensor ID (from DebugMode's _tensor_memo) to the - # capture key of the op that produced it. Populated as we - # walk operators in dispatch order, so by the time we look up - # op N's input IDs the map already contains its predecessors' - # output IDs. Used to attach ``input_producers`` per capture - # so the HTML can show "this input came from " on hover. - producer_map: dict[int, str] = {} - - if self._debug_mode is None: - return - for op in self._debug_mode.operators: - if not isinstance(op, _OpCall): - continue - - op_name = _get_op_name(op.op) - record = op.record or {} - stats = record.get("stats") or {} - - # log_tensor_hashes stamps op.log with input/output norm - # hashes — populated for every _OpCall DebugMode sees, - # including in-place ops where result is None. - log = op.log or {} - output_hash = _stringify_hash_tree(log.get("hash")) - input_hashes = _stringify_hash_tree(log.get("input_hash")) - - # Two cases that produce a capture: - # 1. We have pre-computed stats → normal capture. - # 2. No stats but we have hashes AND the op name ends with - # "_" (in-place convention) and isn't excluded → record - # a hash-only capture so optimizer steps / param updates - # surface in the diff. - is_inplace_capture = ( - not stats - and op_name not in _EXCLUDED_OPS - and op_name.endswith("_") - and (output_hash or input_hashes) - ) - if not stats and not is_inplace_capture: - continue - - module_name = record.get("module_fqn") or "" - seq = op_counters.setdefault(module_name, 0) - op_counters[module_name] = seq + 1 - key = f"{module_name}/op_{seq}_{op_name}" - - stack_frames = record.get("stack_frames") - if not stack_frames: - # Fall back to DebugMode's own traceback strings. - # fwd_stack_trace is the autograd-saved forward stack - # (useful on backward ops); stack_trace is the live - # dispatch-time stack (poor under FSDP / replay). - trace_str = getattr(op, "fwd_stack_trace", None) or getattr( - op, "stack_trace", None - ) - stack_frames = _parse_stack_trace(trace_str) if trace_str else [] - - # Resolve each input tensor ID to the capture key that - # produced it (if any). Positional: empty string means - # "produced by an uncaptured op or external input - # (parameter, dataloader)". - input_ids = record.get("input_ids") or [] - input_producers = ";".join(producer_map.get(tid, "") for tid in input_ids) - - self._captures[key] = CapturedActivation( - stats=stats, - stack_frames=stack_frames, - phase=record.get("phase", "forward"), - output_hash=output_hash, - input_hashes=input_hashes, - input_producers=input_producers, - op_name=op_name, - ) - - # Register this op as the producer of its output tensors so - # downstream consumers can find it. - for tid in record.get("output_ids") or []: - producer_map[tid] = key - - -class ActivationCaptureProfiler: - """Profiler-driven companion that captures per-op activations on a designated step. - - Designed to be driven by :class:`Profiler` via its ``step()`` method, - which is called *after* each training step. To capture step N this - profiler arms ``DebugModeTracer`` after step N-1 completes and dumps - after step N completes. - - Args: - enabled: Whether numerics capture is active. - model: The model whose activations to capture. - dump_dir: Directory for output log files. - capture_step: The training step to capture (1-indexed). - """ - - def __init__( - self, - *, - enabled: bool, - model: nn.Module, - dump_dir: str, - capture_step: int = 1, - ): - self._enabled = enabled - self._model = model - self._dump_dir = dump_dir - self._capture_step = capture_step - self._step = 0 - self._captured = False - self._tracer: DebugModeTracer | None = None - self._captures: dict[str, CapturedActivation] | None = None - - def __enter__(self) -> "ActivationCaptureProfiler": - if self._enabled and self._capture_step == 1: - self._setup() - return self - - def __exit__(self, *args) -> None: - self._teardown() - - def step(self) -> None: - """Called after each training step (same cadence as Profiler.step).""" - self._step += 1 - if not self._enabled or self._captured: - return - if self._step == self._capture_step - 1: - self._setup() - elif self._step == self._capture_step: - self._dump() - - def _setup(self) -> None: - """Enter DebugModeTracer so the next training step is captured.""" - from torchtitan.tools.logging import logger - - logger.info(f"Numerics capture: arming for step {self._capture_step}") - set_numerics_capture_active(True) - self._tracer = DebugModeTracer(self._model) - self._captures = self._tracer.__enter__() - - def _dump(self) -> None: - """Dump captures after the capture step completes.""" - from torchtitan.tools.logging import logger - - # _teardown() exits the tracer. DebugModeTracer populates - # skipped_excluded_ops inside __exit__ (when operators are - # walked), so we snapshot it *after* teardown. - tracer = self._tracer - self._teardown() - skipped = set(tracer.skipped_excluded_ops) if tracer is not None else set() - if not self._captures: - return - - rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 - filepath = os.path.join(self._dump_dir, f"rank_{rank}_activations.log") - dump_captures_to_file(self._captures, filepath, skipped_excluded_ops=skipped) - logger.info(f"Dumped {len(self._captures)} activations to {filepath}") - self._captured = True - self._captures = None - - def _teardown(self) -> None: - set_numerics_capture_active(False) - tracer = self._tracer - self._tracer = None - if tracer is None: - return - if tracer._debug_mode is not None: - tracer.__exit__(None, None, None) diff --git a/agent_tooling/numerics_debugging/compare_numerics.py b/agent_tooling/numerics_debugging/compare_numerics.py deleted file mode 100644 index 8b969f4a11c..00000000000 --- a/agent_tooling/numerics_debugging/compare_numerics.py +++ /dev/null @@ -1,1151 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Compare two numerics activation logs and generate an interactive HTML diff. - -Inputs are the per-rank text logs produced by -``agent_tooling.numerics_debugging.activation_tracer.dump_captures_to_file`` (typically -written to ``{dump_folder}/numerics/rank_{rank}_activations.log`` when -``--profiler.dump_numerics`` is set). Each log is a sequence of -``[module_fqn/op_N_opname]`` blocks with per-op stats (Shape, output hash, -input hashes, L2 norm, Min/Max/Mean), source ``Location``, and optional -``Phase``. - -Pipeline: - 1. ``parse_log`` reads each file into a list of :class:`OpEntry`. - 2. ``match_entries`` aligns the two lists with four passes (manual - overrides → exact key → fuzzy key → numeric stats) so that ops with - different counters or FQNs across modes can still be compared. - 3. ``_compute_diffs`` flags per-stat differences on each stat's own - scale via symmetric relative difference. - 4. ``generate_html`` renders an interactive page with collapsible - Forward/Backward sections, per-cell diff coloring, filter toggles, - and row hover-highlighting. - -The naming "eager" vs "traced" is purely conventional — the tool is -symmetric, and any two logs (e.g. two ranks, two compile modes, two -checkpoints) can be compared. - -Usage: - python compare_numerics.py -o -""" - -import argparse -import html -import math -import re -from dataclasses import dataclass, field - - -@dataclass -class OpEntry: - """One captured op parsed from a numerics log file. - - Attributes: - key: ``module_fqn/op_N_opname`` identifier from the ``[...]`` - header line. Used as the primary matching key. - stats: Per-stat string values (Shape, L2 norm, Min, Max, Mean, and - optionally legacy L1 norm). Stored as strings so we can preserve - formatting and detect non-numeric mismatches (e.g. shape). - phase: ``"forward"`` or ``"backward"`` from the ``Phase:`` line. - Forward is the default when the log omits the line. - location: Source location like ``common/attention.py:534`` from - the ``Location:`` line. Surfaced in the HTML for navigation. - """ - - key: str - stats: dict[str, str] = field(default_factory=dict) - phase: str = "forward" - location: str = "" - # DebugMode norm hash of the op result. Empty for in-place ops - # (no result) and for logs that predate the hash fields. Stored - # as a single string — comma-separated when the op returns a tuple. - output_hash: str = "" - # Comma-separated norm hashes of input tensors, captured *before* - # the op runs. Useful for detecting in-place mutations: comparing - # an op's ``input_hashes`` against the previous op's - # ``input_hashes`` for the same tensor surfaces the change even - # when ``output_hash`` is empty. - input_hashes: str = "" - # Semicolon-separated producing capture key per input tensor, - # aligned positionally with ``input_hashes``. Empty positions - # mean the input came from outside the captured op stream - # (uncaptured intermediate, parameter, dataloader). Used to render - # "produced by " tooltips in the HTML. - input_producers: str = "" - - -def parse_log(path: str) -> tuple[list[OpEntry], set[str]]: - """Parse an activation log file. - - The file format is the one produced by - ``activation_tracer.dump_captures_to_file``: a small header (total - count, optional ``Excluded ops dispatched:`` line) followed by a - sequence of op blocks. Each op block begins with a ``[key]`` - header on its own line, followed by indented stat lines - (``Shape:``, ``L1 norm:``, etc.), and optionally ``Location:`` and - ``Phase:`` lines. Anything else is ignored. - - Args: - path: Path to the log file (typically - ``{dump_folder}/numerics/rank_{rank}_activations.log``). - - Returns: - Tuple of ``(entries, skipped_excluded_ops)``. - Entries are in the order they appear in the file (the - op-dispatch order during the captured step, which is meaningful - for the matching passes in ``match_entries``). - ``skipped_excluded_ops`` is the set of op names that dispatched - but were dropped during capture because they live in - ``_EXCLUDED_OPS``; empty if the log doesn't carry the header - line. - """ - entries: list[OpEntry] = [] - skipped: set[str] = set() - current: OpEntry | None = None - with open(path) as f: - for line in f: - m = re.match(r"^\[(.+)\]", line) - if m: - if current: - entries.append(current) - current = OpEntry(key=m.group(1)) - elif current: - m2 = re.match(r"^\s+(Shape|L1 norm|L2 norm|Min|Max|Mean):\s+(.+)", line) - if m2: - current.stats[m2.group(1)] = m2.group(2).strip() - elif "Location:" in line: - current.location = line.strip().replace("Location: ", "") - elif "Phase:" in line: - current.phase = line.strip().replace("Phase: ", "") - elif "Output hash:" in line: - current.output_hash = line.strip().replace("Output hash: ", "") - elif "Input hashes:" in line: - current.input_hashes = line.strip().replace("Input hashes: ", "") - elif "Input producers:" in line: - current.input_producers = line.strip().replace( - "Input producers: ", "" - ) - elif line.startswith("Excluded ops dispatched:"): - payload = line.split(":", 1)[1].strip() - if payload and payload != "(none)": - skipped = {op.strip() for op in payload.split(",") if op.strip()} - if current: - entries.append(current) - return entries, skipped - - -# Leading CamelCase class-name segment in a key, e.g. -# ``FSDPLlama3Model.layers.0...``. ModTracker (used by the DebugMode -# eager backend) roots FQNs at the wrapped model class, while -# FQNInterpreter (used during traced replay) uses the unprefixed -# ``layers.0...`` form. The pattern requires the segment to start with -# an uppercase letter and contain another uppercase letter, so common -# all-lowercase attribute names (``layers``, ``feed_forward``) are -# never accidentally stripped. -_ROOT_CLASS_PREFIX_RE = re.compile(r"^[A-Z][A-Za-z0-9_]*[A-Z][A-Za-z0-9_]*\.") - - -def _strip_root_class_prefix(key: str) -> str: - """Remove a leading ``ClassName.`` segment from an FQN. - - Backend-agnostic match normalization: the DebugMode eager path - produces ``FSDPLlama3Model.layers.0.attention.wo/op_0_mm`` while - the traced path produces ``layers.0.attention.wo/op_0_mm``. - Stripping the class segment lets them match by key. - - Generalizes across models — ``FSDPQwenModel.``, ``Llama2Model.``, - etc. are all matched by the same heuristic. - """ - return _ROOT_CLASS_PREFIX_RE.sub("", key, count=1) - - -def _fuzzy_key(key: str) -> str: - """Strip the per-module op counter from a key. - - Example: ``layers.0.w2/op_3_mm`` -> ``layers.0.w2/mm``. - - Used by ``match_entries`` Pass 2: between modes the same module may - dispatch the same op N times but in different relative order (e.g. - backward ops are interleaved differently), so the counter ``N`` is - unstable. Stripping it lets us match by ``module_fqn + op_type``. - - Also strips a leading ``ClassName.`` segment so eager DebugMode - keys (rooted at the model class) match traced keys (unrooted). - """ - return _strip_root_class_prefix(re.sub(r"/op_\d+_", "/", key)) - - -# Stat names recognized in the log and compared per-op. Must match the -# stats written by activation_tracer.dump_captures_to_file. Adding a new -# stat here also requires extending the parser regex in parse_log and -# the table headers in generate_html. L1 norm is intentionally absent — -# it duplicates ``output_hash`` (the DebugMode norm hash is L1 in -# float64), so we render that column as "Output L1 norm" instead. -STAT_FIELDS = ["Shape", "L2 norm", "Min", "Max", "Mean"] - -# Column-header labels for the rendered HTML. Keeping the underlying -# OpEntry.stats keys short (matching the log) while exposing -# user-friendly headers that make it clear these stats describe the -# *output* tensor. Shape is left bare since it's the only one not -# obviously an aggregate. -_STAT_DISPLAY_LABELS = { - "Shape": "Shape", - "L2 norm": "Output L2 norm", - "Min": "Output Min", - "Max": "Output Max", - "Mean": "Output Mean", -} - - -def _rel_diff(a_str: str, b_str: str) -> float | None: - """Compute symmetric relative difference between two stat strings. - - The denominator is ``max(|a|, |b|, 1e-30)`` so the result is bounded - by 1 and never blows up when both values are tiny. The 1e-30 floor - prevents division by zero when both inputs are exactly zero. - - Args: - a_str: First stat string (e.g. ``"3.36e+04"``). - b_str: Second stat string. - - Returns: - Relative difference in ``[0, 1]`` for numeric inputs, or ``None`` - when either value cannot be parsed as a float (e.g. the - ``Shape`` field, which is rendered as ``torch.Size([...])``). - """ - try: - a, b = float(a_str), float(b_str) - except (ValueError, TypeError): - return None - # Treat NaN inputs as "non-numeric" so the caller falls back to - # exact-string equality. Async-collective outputs hash to NaN - # before sync; comparing NaN propagates through the relative diff - # and breaks the downstream intensity gradient. - if math.isnan(a) or math.isnan(b): - return None - denom = max(abs(a), abs(b), 1e-30) - return abs(a - b) / denom - - -def _hash_intensity(e_val: str, t_val: str) -> float | None: - """Max relative diff between two comma-separated hash lists. - - Returns: - - ``None`` when both sides are empty. - - ``0.0`` when the strings are identical or numerically equal. - - ``1.0`` for structural mismatches (different list length, or - one side empty, or non-numeric tokens that don't match). - - The largest position-wise relative diff (``_rel_diff``) - otherwise. - - Used for both the HTML cell coloring intensity and the - ``_compute_match_status`` threshold. - """ - if not (e_val or t_val): - return None - if not e_val or not t_val: - return 1.0 - if e_val == t_val: - return 0.0 - ev = [v.strip() for v in e_val.split(",")] - tv = [v.strip() for v in t_val.split(",")] - if len(ev) != len(tv): - return 1.0 - max_rd = 0.0 - for a_s, b_s in zip(ev, tv): - # ``nan`` shows up for async-collective outputs that haven't - # synced yet (the destination buffer is uninitialized). Two - # NaNs side-by-side = both sides observed the same unsynced - # bucket → treat as matching. One side NaN + one numeric = - # structural mismatch. - a_nan = a_s == "nan" - b_nan = b_s == "nan" - if a_nan or b_nan: - if a_nan and b_nan: - continue - return 1.0 - rd = _rel_diff(a_s, b_s) - if rd is None: - if a_s != b_s: - return 1.0 - continue - if rd > max_rd: - max_rd = rd - return max_rd - - -def _compute_match_status(e_entry, t_entry) -> str: - """Classify a paired row as ``"match"`` or ``"diff"``. - - The classification considers only ``Shape``, ``output_hash``, and - ``input_hashes`` — the three fields that uniquely identify "this - is the same activation". L2 norm / Min / Max / Mean are - sub-statistics of the same tensor; they get displayed (with - intensity coloring) for context, but they don't drive the toggle. - - Threshold: hash relative-diff > 1e-8 counts as a real difference, - matching the floor used by the hash-cell color gradient (smaller - diffs are pure float64 reduction noise). - """ - if e_entry.stats.get("Shape", "") != t_entry.stats.get("Shape", ""): - return "diff" - for e_val, t_val in ( - (e_entry.output_hash, t_entry.output_hash), - (e_entry.input_hashes, t_entry.input_hashes), - ): - rd = _hash_intensity(e_val, t_val) - if rd is not None and rd > 1e-8: - return "diff" - return "match" - - -def _compute_diffs(e_entry, t_entry): - """Compute per-stat differences between two matched op entries. - - Each stat is compared on its own scale via :func:`_rel_diff` - (symmetric relative difference, ``abs(a-b) / max(|a|, |b|, 1e-30)``). - L1 norm differences are judged against L1 magnitude, Min against - Min magnitude, etc. — comparing across columns (e.g. scaling Min - by L1) would over- or under-suppress diffs depending on stat. - - Special case: non-numeric stats (Shape, rendered as - ``torch.Size([...])``) get a flat ``1.0`` intensity when they - differ, since "the shape changed" is binary. - - Args: - e_entry: Entry from the first log (typically eager). - t_entry: Entry from the second log (typically traced). - - Returns: - Dict mapping stat name to a relative-diff intensity in - ``[0, 1]`` (numeric stats) or ``1.0`` (non-numeric mismatch). - Only stats that actually differ are included; an empty dict - means "matched". - """ - diffs = {} - for stat in STAT_FIELDS: - ev = e_entry.stats.get(stat) - tv = t_entry.stats.get(stat) - if ev and tv and ev != tv: - rd = _rel_diff(ev, tv) - # _rel_diff returns None for non-numeric stats (Shape). - diffs[stat] = 1.0 if rd is None else rd - return diffs - - -def match_entries( - eager: list[OpEntry], - traced: list[OpEntry], - overrides: dict[str, str] | None = None, -): - """Pair up ops between two logs using a four-pass matching strategy. - - Cross-mode comparisons (eager vs aot_fx_trace, etc.) - can't rely on identical keys: op counters drift when ops are - interleaved differently, and some ops have different module - attributions between modes (e.g. gradient accumulation adds may - surface under different FQNs). The passes progressively relax the - matching criteria so that ops still get paired: - - Pass 0 — Manual overrides (only if ``overrides`` is provided). - Force-pairs the named keys regardless of stats / signatures. - Use this to fix-up specific cases where the FQN scheme drifts - between runs in a way the automated passes can't recover (e.g. - an AC-recomputed op surfacing as bare ``feed_forward/op_7_mul`` - in eager but ``layers.2.feed_forward/op_3_mul`` in traced). - - Pass 1 — Exact key match. Catches the bulk of forward ops where - both runs dispatch the same op in the same module-relative - order. Most-specific match; runs first to claim the - "obviously the same op" cases. - - Pass 2 — Fuzzy key match (``module_fqn + op_type``, counter - stripped via :func:`_fuzzy_key`). Handles ops where the - per-module sequence number drifts between modes, common for - backward ops which are interleaved differently. - - Pass 3 — Stats match (``op_type + Shape + L1 norm``). Last-resort - match for ops whose module FQN differs between modes (e.g. - gradient accumulation adds in autograd internals). Identifies - ops by their tensor signature rather than by name. - - Anything still unmatched is returned with status ``eager_only`` or - ``traced_only`` so it shows up in the HTML's "only" filter. - - Args: - eager: Entries from the first log. - traced: Entries from the second log. - overrides: Optional ``{eager_key: traced_key}`` dict. Each - entry forces a pair regardless of stat similarity; any - entry whose target key isn't present on the traced side is - silently ignored (so override files can outlive log - regenerations). - - Returns: - List of ``(eager_entry | None, traced_entry | None, status, - diffs, match_strategy)`` tuples. ``status`` is one of - ``"match"``, ``"diff"``, ``"eager_only"``, ``"traced_only"``. - ``match_strategy`` records which pass paired the entries: - ``"override"``, ``"exact"``, ``"fuzzy"``, ``"stats"``, or - ``""`` for only-side rows. Eager-side entries appear in their - original log order; ``traced_only`` entries follow at the end - in their original traced-log order. - """ - traced_used = set() - overrides = overrides or {} - - # Collect a (status, diffs, traced_entry, strategy) for each eager - # entry, by its index in the eager list. Each pass fills in slots - # that earlier passes left empty. Emitting in eager-index order at - # the end keeps the output aligned with the eager log's op-dispatch - # sequence. - eager_results: list[tuple[str, dict, OpEntry, str] | None] = [None] * len(eager) - - def _try_pass(strategy: str, key_fn) -> None: - # Build traced-side index from currently unclaimed entries, then - # walk eager in order and claim the first available match per - # entry that hasn't been filled by an earlier pass. - index: dict[object, list[OpEntry]] = {} - for e in traced: - if id(e) not in traced_used: - index.setdefault(key_fn(e), []).append(e) - for i, e_entry in enumerate(eager): - if eager_results[i] is not None: - continue - for c in index.get(key_fn(e_entry), []): - if id(c) not in traced_used: - traced_used.add(id(c)) - diffs = _compute_diffs(e_entry, c) - status = _compute_match_status(e_entry, c) - eager_results[i] = (status, diffs, c, strategy) - break - - # Pass 0: manual overrides. Each entry forces a specific - # eager_key -> traced_key pair and consumes both sides so later - # passes don't reassign them. Index traced by raw key for lookup. - if overrides: - traced_by_key: dict[str, list[OpEntry]] = {} - for t in traced: - traced_by_key.setdefault(t.key, []).append(t) - for i, e_entry in enumerate(eager): - target_key = overrides.get(e_entry.key) - if target_key is None: - continue - for c in traced_by_key.get(target_key, []): - if id(c) not in traced_used: - traced_used.add(id(c)) - diffs = _compute_diffs(e_entry, c) - status = _compute_match_status(e_entry, c) - eager_results[i] = (status, diffs, c, "override") - break - - # Pass 1: exact match (forward ops, where keys agree). The root - # class prefix is normalized here too — DebugMode eager keys are - # rooted at the model class (``FSDPLlama3Model.``), traced keys - # are not, so a literal e.key equality would never match across - # backends. - _try_pass("exact", lambda e: _strip_root_class_prefix(e.key)) - - # Pass 2: fuzzy match (module_fqn + op_type, counter stripped). - _try_pass("fuzzy", lambda e: _fuzzy_key(e.key)) - - # Pass 3: stats match (op_type + Shape + output L1-norm hash) — - # handles ops whose FQN differs between modes. Uses ``output_hash`` - # rather than the now-removed ``L1 norm`` stat field; both encode - # the same float64 L1 reduction, so the key is unchanged in spirit. - def _stats_key(entry: OpEntry) -> tuple[str, str, str]: - op_type = ( - _fuzzy_key(entry.key).split("/")[-1] if "/" in entry.key else entry.key - ) - return (op_type, entry.stats.get("Shape", ""), entry.output_hash) - - _try_pass("stats", _stats_key) - - # Emit eager entries in their original order. Anything still None - # is eager_only. Typed as ``list[tuple]`` so pyrefly doesn't narrow - # element type to the first append and reject the others. - results: list[tuple] = [] - for e_entry, slot in zip(eager, eager_results, strict=True): - if slot is None: - results.append((e_entry, None, "eager_only", {}, "")) - else: - status, diffs, t_entry, strategy = slot - results.append((e_entry, t_entry, status, diffs, strategy)) - - # traced_only entries (never claimed by any pass) trail at the end - # in traced-log order. - for t_entry in traced: - if id(t_entry) not in traced_used: - results.append((None, t_entry, "traced_only", {}, "")) - - return results - - -def generate_html( - results, - eager_path, - traced_path, - name1: str = "run1", - name2: str = "run2", - skipped1: set[str] | None = None, - skipped2: set[str] | None = None, -) -> str: - """Render matching results as a self-contained interactive HTML page. - - The output has: - - * A summary bar with match / diff / only-side counts. - * Filter checkboxes to hide each row category. - * Two collapsible sections — Forward and Backward — each rendering - a table with one row per (run1, run2) pair. Each row shows both - sides' stats stacked, with diff cells colored on a white-to-red - log scale, and the source ``Location`` for click-to-navigate - context. - * Hover-highlighting that tints both halves of a row pair together. - - The page is fully self-contained (inline CSS + JS, no external - assets). - - Args: - results: Output of :func:`match_entries`. - eager_path: Path of the first log (currently unused in output; - kept so callers can pass it for future provenance display). - traced_path: Path of the second log. - name1: Display name for the first log (defaults to ``"run1"``). - Used in the page title, row-side label, and filter - checkbox. - name2: Display name for the second log (defaults to ``"run2"``). - skipped1: Op names that dispatched in run 1 but were dropped by - ``_EXCLUDED_OPS`` (from the log's ``Excluded ops - dispatched:`` header). Rendered as a small chip list under - the summary so the user sees which always-skipped ops the - model actually exercises. - skipped2: Same for run 2. - - Returns: - Complete HTML document as a string. - """ - # HTML-escape user-provided names since they go into the document - # body and attribute values. - name1_safe = html.escape(name1) - name2_safe = html.escape(name2) - - def _skipped_block(label: str, ops: set[str] | None) -> str: - if not ops: - chips = '(none)' - else: - chips = "".join( - f'{html.escape(op)}' - for op in sorted(ops) - ) - return ( - f'
' - f'{html.escape(label)}: {chips}' - f"
" - ) - - skipped_section = ( - f'
' - f'
Excluded ops dispatched ' - f'(present at runtime but filtered ' - f"by _EXCLUDED_OPS)
" - f"{_skipped_block(name1, skipped1)}" - f"{_skipped_block(name2, skipped2)}" - f"
" - ) - - def _phase(e, t): - return e.phase if e else t.phase - - fwd = [r for r in results if _phase(r[0], r[1]) == "forward"] - bwd = [r for r in results if _phase(r[0], r[1]) == "backward"] - - total_match = sum(1 for r in results if r[2] == "match") - total_diff = sum(1 for r in results if r[2] == "diff") - total_eager = sum(1 for r in results if r[2] == "eager_only") - total_traced = sum(1 for r in results if r[2] == "traced_only") - - def make_table(entries, section_id): - rows: list[str] = [] - for idx, (e, t, status, diffs, strategy) in enumerate(entries): - row_id = f"{section_id}_{idx}" - - if status == "eager_only": - cls = "eager-only" - elif status == "traced_only": - cls = "traced-only" - elif status == "diff": - cls = "has-diff" - else: - cls = "all-match" - - # Show both keys when they differ (fuzzy match) - e_key = html.escape(e.key) if e else "" - t_key = html.escape(t.key) if t else "" - if e and t and e.key == t.key: - key_html = e_key - elif e and t: - key_html = ( - f'{e_key}' - f'
{t_key}' - ) - else: - key_html = e_key or t_key - - # Build stat cells with intensity-based diff coloring. - # Relative diff is mapped to opacity: tiny diffs are faint - # red, large diffs and non-numeric mismatches (Shape) are - # full red. - eager_cells: list[str] = [] - traced_cells: list[str] = [] - for stat in STAT_FIELDS: - ev = e.stats.get(stat, "") if e else "" - tv = t.stats.get(stat, "") if t else "" - rd = diffs.get(stat) - if rd is not None: - # White-to-red gradient: intensity 0 = white text, - # intensity 1 = full red text + red background tint. - # Log scale: 1e-8 → 0.0, 1e-1+ → 1.0 - if rd >= 0.1: - t_val = 1.0 - elif rd <= 1e-8: - t_val = 0.0 - else: - t_val = (math.log10(rd) + 8) / 7 - # Text: white (220,220,220) → red (248,80,80) - r = int(220 + (248 - 220) * t_val) - g = int(220 - (220 - 80) * t_val) - b = int(220 - (220 - 80) * t_val) - # Background: transparent → faint red - bg_alpha = t_val * 0.15 - style = ( - f' style="color: rgb({r},{g},{b});' - f" background: rgba(248,80,80,{bg_alpha:.2f});" - f' font-weight: bold"' - ) - else: - style = "" - # Shape is rendered as ``torch.Size([...]), Dtype: ...`` - # which can be long for high-rank tensors; wrap in the - # same collapsible cell used by the L1-norm columns. - if stat == "Shape": - ev_inner = ( - f'" - if ev - else "" - ) - tv_inner = ( - f'" - if tv - else "" - ) - else: - ev_inner = html.escape(ev) - tv_inner = html.escape(tv) - eager_cells.append(f"{ev_inner}") - traced_cells.append(f"{tv_inner}") - - e_loc = html.escape(e.location) if e else "" - t_loc = html.escape(t.location) if t else "" - e_phase = e.phase if e else "" - t_phase = t.phase if t else "" - - # Hash cells: input and output norm hashes from - # log_tensor_hashes. Comma-separated lists of floats — we - # compare element-wise with the same log-scale relative- - # diff intensity used for the L1/L2/Min/Max columns, so a - # last-ULP drift (~1e-7) renders as nearly invisible faint - # pink and a true divergence (~1e-1) renders as full red. - # A structural mismatch (different list length) renders as - # full red — that genuinely means the op signatures differ. - # - # For input hashes we additionally wrap each individual - # value in a when a producer key is - # available — hovering shows "produced by ". - def _format_hash_values(val_str: str, prod_str: str) -> str: - if not val_str: - return "" - vals = [v.strip() for v in val_str.split(",")] - prods = [p.strip() for p in prod_str.split(";")] if prod_str else [] - parts: list[str] = [] - for i, v in enumerate(vals): - prod = prods[i] if i < len(prods) else "" - if prod: - # Both the CSS tooltip (data-tooltip via ::after) - # and native title are emitted — title is the - # universal fallback if CSS gets stripped by the - # host environment. - tip = html.escape(f"produced by {prod}") - parts.append( - f'' - f"{html.escape(v)}" - ) - else: - parts.append(html.escape(v)) - return ", ".join(parts) - - def _hash_cell( - e_val: str, - t_val: str, - e_prod: str = "", - t_prod: str = "", - collapsible: bool = False, - ): - e_inner = _format_hash_values(e_val, e_prod) - t_inner = _format_hash_values(t_val, t_prod) - if collapsible: - # Wrap in a clickable div that starts collapsed - # (max-height limited) and expands on click. Lets - # rows with hundreds of input hashes (e.g. - # _fused_adamw_) stay readable. - def _wrap(inner: str) -> str: - if not inner: - return "" - return ( - '" - ) - - e_inner = _wrap(e_inner) - t_inner = _wrap(t_inner) - if not (e and t): - return f"{e_inner}", f"{t_inner}" - rd = _hash_intensity(e_val, t_val) - if rd is None or rd == 0.0: - return f"{e_inner}", f"{t_inner}" - if rd >= 0.1: - t_val_intensity = 1.0 - elif rd <= 1e-8: - t_val_intensity = 0.0 - else: - t_val_intensity = (math.log10(rd) + 8) / 7 - r = int(220 + (248 - 220) * t_val_intensity) - g = int(220 - (220 - 80) * t_val_intensity) - b = int(220 - (220 - 80) * t_val_intensity) - bg_alpha = t_val_intensity * 0.15 - style = ( - f' style="color: rgb({r},{g},{b});' - f" background: rgba(248,80,80,{bg_alpha:.2f});" - f' font-weight: bold"' - ) - return f"{e_inner}", f"{t_inner}" - - e_oh = e.output_hash if e else "" - t_oh = t.output_hash if t else "" - e_oh_cell, t_oh_cell = _hash_cell(e_oh, t_oh, collapsible=True) - e_ih = e.input_hashes if e else "" - t_ih = t.input_hashes if t else "" - e_ip = e.input_producers if e else "" - t_ip = t.input_producers if t else "" - e_ih_cell, t_ih_cell = _hash_cell(e_ih, t_ih, e_ip, t_ip, collapsible=True) - - strategy_label = { - "override": "manual override", - "exact": "same op key", - "fuzzy": "fuzzy op key", - "stats": "stats", - }.get(strategy, "") - strategy_html = ( - f'{strategy_label}' - if strategy - else "" - ) - - rows.append( - f""" - - {key_html} - {strategy_html} - {name1_safe} - {e_oh_cell} - {e_ih_cell} - {e_loc} - {eager_cells[0]} - {e_phase} - {''.join(eager_cells[1:])} - - - {name2_safe} - {t_oh_cell} - {t_ih_cell} - {t_loc} - {traced_cells[0]} - {t_phase} - {''.join(traced_cells[1:])} - """ - ) - - return "\n".join(rows) - - return f""" - - - -Numerics Comparison - - - -

Numerics Comparison: {name1_safe} vs {name2_safe}

-
- ✓ {total_match} matched - ✗ {total_diff} diffs - ← {total_eager} {name1_safe} only - → {total_traced} {name2_safe} only -
- -{skipped_section} - -
- - - - - - -
- -
-

Forward ({len(fwd)} ops) ▼

-
- - - - - - - - - {make_table(fwd, "fwd")} -
OpMatch methodSideOutput L1 normInput L1 normLocationShapePhaseOutput L2 normOutput MinOutput MaxOutput Mean
-
-
- -
-

Backward ({len(bwd)} ops) ▼

-
- - - - - - - - - {make_table(bwd, "bwd")} -
OpMatch methodSideOutput L1 normInput L1 normLocationShapePhaseOutput L2 normOutput MinOutput MaxOutput Mean
-
-
- - - - -""" - - -def main(): - """CLI entry point: parse two logs, match entries, write HTML diff. - - Invoked via ``python -m agent_tooling.numerics_debugging.compare_numerics`` or - directly. See module docstring for the typical workflow. - """ - parser = argparse.ArgumentParser(description="Compare numerics activation logs") - parser.add_argument("eager", help="First log file") - parser.add_argument("traced", help="Second log file") - parser.add_argument( - "-o", "--output", default="numerics_diff.html", help="Output HTML file" - ) - parser.add_argument( - "--name1", - default="run1", - help="Display name for the first log (default: run1)", - ) - parser.add_argument( - "--name2", - default="run2", - help="Display name for the second log (default: run2)", - ) - parser.add_argument( - "--override", - default=None, - help=( - "Path to a CSV of forced matches: each non-empty line is " - "', '. Applied before the automated " - "matching passes; matching pairs are rendered with the " - "'manual override' strategy chip. Lines starting with '#' " - "are ignored." - ), - ) - args = parser.parse_args() - - overrides: dict[str, str] = {} - if args.override: - with open(args.override) as f: - for raw in f: - line = raw.strip() - if not line or line.startswith("#"): - continue - parts = line.split(",", 1) - if len(parts) != 2: - continue - overrides[parts[0].strip()] = parts[1].strip() - print(f"Loaded {len(overrides)} override(s) from {args.override}") - - eager, skipped_eager = parse_log(args.eager) - traced, skipped_traced = parse_log(args.traced) - results = match_entries(eager, traced, overrides=overrides) - html_content = generate_html( - results, - args.eager, - args.traced, - name1=args.name1, - name2=args.name2, - skipped1=skipped_eager, - skipped2=skipped_traced, - ) - - with open(args.output, "w") as f: - f.write(html_content) - print(f"Wrote {args.output} ({len(results)} entries)") - - -if __name__ == "__main__": - main() diff --git a/assets/version.txt b/assets/version.txt index ee1372d33a2..35c68595328 100644 --- a/assets/version.txt +++ b/assets/version.txt @@ -1 +1 @@ -0.2.2 +0.2.2+neuron.1 diff --git a/pyproject.toml b/pyproject.toml index 9f9f60fc7d3..df3970c928d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,11 @@ # ---- All project specifications ---- # [project] -name = "torchtitan" +name = "torchtitan-neuron" description = "A PyTorch native platform for training generative AI models" readme = "README.md" requires-python = ">=3.10" license = {file = "LICENSE"} -authors = [ - { name = "PyTorch Team", email = "packages@pytorch.org" }, -] +authors = [{name = "AWS Neuron Team"}] keywords = ["pytorch", "training", "llm"] dependencies = [ # Stateful Dataloader @@ -29,9 +27,10 @@ dependencies = [ dynamic = ["version"] [project.urls] -GitHub = "https://github.com/pytorch/torchtitan" -Documentation = "https://github.com/pytorch/torchtitan/tree/main/docs" -Issues = "https://github.com/pytorch/torchtitan/issues" +GitHub = "https://github.com/aws-neuron/torchtitan-neuron" +Documentation = "https://github.com/aws-neuron/torchtitan-neuron/tree/main/docs" +Issues = "https://github.com/aws-neuron/torchtitan-neuron/issues" +Upstream = "https://github.com/pytorch/torchtitan" [project.optional-dependencies] dev = [ @@ -57,10 +56,14 @@ build-backend = "setuptools.build_meta" [tool.setuptools.packages.find] where = [""] include = ["torchtitan*"] +exclude = ["*.tests", "*.tests.*"] + +[tool.setuptools.package-data] +"torchtitan" = ["assets/**/*"] [tool.pytest.ini_options] addopts = ["--showlocals"] # show local variables in tracebacks -testpaths = ["tests"] +testpaths = ["tests", "torchtitan_neuron/tests"] [tool.pyrefly] project-excludes = ["torchtitan/experiments", "**/tests/**"] diff --git a/run_train.sh b/run_train.sh index d2940d00817..f19f66ff048 100755 --- a/run_train.sh +++ b/run_train.sh @@ -1,9 +1,10 @@ #!/usr/bin/bash # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. - -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. set -ex @@ -26,22 +27,37 @@ set -ex # Example: NGPU=32 COMM_MODE="local_tensor" ./run_train.sh NGPU=${NGPU:-"8"} +NNODES=${NNODES:-"1"} +RDZV_ENDPOINT=${RDZV_ENDPOINT:-"localhost:0"} export LOG_RANK=${LOG_RANK:-0} MODULE=${MODULE:-"llama3"} CONFIG=${CONFIG:-"llama3_debugmodel"} COMM_MODE=${COMM_MODE:-""} +# An optional first .sh argument can prepare the launch environment. +# Example: ./run_train.sh path/to/environment.sh --training.steps 10 +if [[ "${1:-}" == *.sh ]]; then + ENVIRONMENT_SCRIPT=$1 + if [[ ! -f "${ENVIRONMENT_SCRIPT}" ]]; then + echo "Environment script does not exist: ${ENVIRONMENT_SCRIPT}" >&2 + exit 1 + fi + source "${ENVIRONMENT_SCRIPT}" + shift +fi + TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE:-"http://localhost:29510"} if [ -n "$COMM_MODE" ]; then # Communication mode specified: validate configuration or run in debug mode echo "Running with comm_mode=${COMM_MODE}" - NGPU="${NGPU}" LOCAL_RANK=0 python3 -m torchtitan.train --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 + NGPU="${NGPU}" LOCAL_RANK=0 python3 -m torchtitan.train --module "${MODULE}" --config "${CONFIG}" "$@" --comm.mode="${COMM_MODE}" --training.steps 1 else # Normal training with torchrun PYTORCH_ALLOC_CONF="expandable_segments:True" \ - TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ - torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ - --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ - -m torchtitan.train --module ${MODULE} --config ${CONFIG} "$@" + TORCHFT_LIGHTHOUSE="${TORCHFT_LIGHTHOUSE}" \ + torchrun --nnodes="${NNODES}" --nproc_per_node="${NGPU}" \ + --rdzv_backend c10d --rdzv_endpoint="${RDZV_ENDPOINT}" \ + --local-ranks-filter "${LOG_RANK}" --role rank --tee 3 \ + -m torchtitan.train --module "${MODULE}" --config "${CONFIG}" "$@" fi diff --git a/tests/README.md b/tests/README.md index 3e0e27b106f..3a2a2448ff4 100644 --- a/tests/README.md +++ b/tests/README.md @@ -64,6 +64,14 @@ To run only the unit tests: pytest -s tests/unit_tests/ ``` +Core attention backend selection and model registry contracts are covered by +`tests/unit_tests/test_attn_backend.py`. These tests run on CPU and are included +in the CPU unit-test CI job: + +```bash +pytest -s tests/unit_tests/test_attn_backend.py +``` + ### Running Specific Unit Test Files To run a specific test file: diff --git a/tests/integration_tests/__init__.py b/tests/integration_tests/__init__.py index 6fb01ceb541..dba8138b8d1 100644 --- a/tests/integration_tests/__init__.py +++ b/tests/integration_tests/__init__.py @@ -1,16 +1,21 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. +import importlib.util from collections.abc import Sequence from dataclasses import dataclass __all__ = [ "OverrideDefinitions", + "_ON_NEURON", ] +_ON_NEURON = importlib.util.find_spec("torch_neuronx") is not None + @dataclass class OverrideDefinitions: @@ -24,6 +29,7 @@ class OverrideDefinitions: ngpu: int = 4 disabled: bool = False skip_rocm_test: bool = False + skip_neuron: bool = False timeout: int | None = None def __repr__(self): diff --git a/tests/integration_tests/features.py b/tests/integration_tests/features.py index 179176edcfc..315bf755f1b 100755 --- a/tests/integration_tests/features.py +++ b/tests/integration_tests/features.py @@ -1,14 +1,39 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import dataclasses import os -from tests.integration_tests import OverrideDefinitions +from tests.integration_tests import _ON_NEURON, OverrideDefinitions + + +def _inject_neuron_compile_backend( + flavors: list[OverrideDefinitions], +) -> list[OverrideDefinitions]: + """Select the Neuron compile backend without changing model recipes.""" + injected: list[OverrideDefinitions] = [] + for t in flavors: + new_variants = [] + for variant in t.override_args: + args = list(variant) + # Place flags before the activation-checkpoint subcommand. + ac_idx = next( + ( + i + for i, a in enumerate(args) + if a.startswith("activation-checkpoint:") + ), + len(args), + ) + if not any(a.startswith("--compile.backend") for a in args): + args.insert(ac_idx, "--compile.backend neuron") + new_variants.append(tuple(args)) + injected.append(dataclasses.replace(t, override_args=tuple(new_variants))) + return injected def _is_pp_only(variant: tuple[str, ...], ngpu: int) -> bool: @@ -344,18 +369,15 @@ def build_features_test_list() -> list[OverrideDefinitions]: "pp_looped_1f1b", ngpu=4, ), - # TODO: Disabled with the FlexAttention default (SDPA is no longer a - # language-model backend). Zero-bubble / multi schedules split backward - # and call torch's stage_backward_input, which runs - # _get_grad_fn_or_grad_acc (t.requires_grad) over every stage input — + # TODO: Disabled with the default FlexAttention backend. Zero-bubble / + # multi schedules split backward and call torch's stage_backward_input, + # which runs _get_grad_fn_or_grad_acc (t.requires_grad) over every input, # including the forwarded FlexAttention BlockMask, which is not a Tensor # ("'BlockMask' object has no attribute 'requires_grad'"). Full-backward # schedules (1F1B/GPipe/Interleaved1F1B) are unaffected. Re-enable once # stage_backward_input skips non-tensor stage inputs upstream. - # (VarlenAttention's tensor-based metadata would sidestep this, but - # varlen requires flash_attn_interface/FA3, which the core integration - # CI does not install; SDPA is no longer a core LM backend. So the - # upstream stage_backward_input fix is the path here.) + # SDPA avoids BlockMask but does not preserve packed-document masking, + # so it would change the attention semantics exercised by these tests. OverrideDefinitions( [ [ @@ -569,6 +591,7 @@ def build_features_test_list() -> list[OverrideDefinitions]: "Override: swap FeedForward with fused SwiGLU (FSDP2 + TP2)", "override_fused_swiglu", ngpu=4, + skip_neuron=True, ), OverrideDefinitions( [ @@ -583,6 +606,7 @@ def build_features_test_list() -> list[OverrideDefinitions]: "(FSDP2 + TP2 dense, EP4 sparse)", "override_fused_grouped_experts", ngpu=4, + skip_neuron=True, ), # NOTE: below are tests which require config change that cannot be done # via CLI overrides, so remain llama3 specific @@ -598,6 +622,7 @@ def build_features_test_list() -> list[OverrideDefinitions]: "fsdp+varlen_attn+per_op_sac", ngpu=4, skip_rocm_test=True, + skip_neuron=True, ), OverrideDefinitions( [ @@ -671,6 +696,10 @@ def build_features_test_list() -> list[OverrideDefinitions]: ), ] - return [ - *[_enable_spmd_backend(t, "spmd_types") for t in integration_tests_flavors], - ] + # Emit the plain-backend flavors. The spmd_types wrapper + # (``_enable_spmd_backend``) is intentionally not applied: the default + # backend is onboarded, spmd_types is not yet + flavors = list(integration_tests_flavors) + if _ON_NEURON: + flavors = _inject_neuron_compile_backend(flavors) + return flavors diff --git a/tests/integration_tests/run_tests.py b/tests/integration_tests/run_tests.py index cf3d3a4d0c4..019e175e446 100644 --- a/tests/integration_tests/run_tests.py +++ b/tests/integration_tests/run_tests.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import argparse import os @@ -14,7 +15,7 @@ from torchtitan.tools.logging import logger -from tests.integration_tests import OverrideDefinitions +from tests.integration_tests import _ON_NEURON, OverrideDefinitions from tests.integration_tests.features import build_features_test_list from tests.integration_tests.h100 import build_h100_tests_list from tests.integration_tests.models import build_model_tests_list @@ -198,6 +199,8 @@ def _filter_tests( and test_flavor.skip_rocm_test ): continue + if _ON_NEURON and test_flavor.skip_neuron: + continue if args.ngpu < test_flavor.ngpu: skipped_ngpu.append(test_flavor) continue diff --git a/tests/unit_tests/__init__.py b/tests/unit_tests/__init__.py index 2e41cd717f6..2adf0a16f02 100644 --- a/tests/unit_tests/__init__.py +++ b/tests/unit_tests/__init__.py @@ -1,5 +1,50 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron enablement shim for PyTorch's DTensorTestBase. + +This runs at import of the ``tests.unit_tests`` package. It must live here (not +in ``conftest.py``) because ``MultiProcessTestCase`` launches each rank in a +fresh spawned interpreter that never loads ``conftest.py`` -- but every spawned +rank re-imports the test module, and therefore this package ``__init__``. This +is the only place a patch reaches the ranks where the process group is built. + +Problem: torch_neuronx registers ``neuron`` as a PrivateUse1 accelerator with a +``neuron`` c10d backend, so ``DTensorTestBase.backend`` resolves to the bare +string ``"neuron"``. PyTorch's ``common_dtensor.init_pg`` only accepts a fixed +set of backend strings; the bare device backend is rejected with "Backend +neuron not supported!". The compound ``"cpu:gloo,:"`` form IS +in that accepted set, so we return that instead. + +UPSTREAM FIX: recent PyTorch widens ``init_pg``'s accepted-backend list to also +accept a device's registered backend (bare and compound) for any PrivateUse1 +accelerator. Once the torch build under test carries that change, the bare +``"neuron"`` backend is accepted directly and this shim is unnecessary -- it can +be removed (and the enabled tests dropped from the conftest skip list). +""" + +import torch +import torch.distributed as dist +from torch.testing._internal.distributed._tensor import common_dtensor + +# Name of the registered PrivateUse1 backend (e.g. "neuron"); "privateuseone" +# when none is registered. +_PRIVATEUSE1_NAME = torch._C._get_privateuse1_backend_name() + +_original_backend = common_dtensor.DTensorTestBase.backend + + +def _backend(self): + # Only reshape the backend string for the PrivateUse1 accelerator (neuron); + # cpu (gloo) and cuda (nccl) already resolve to an accepted bare backend. + if self.device_type == _PRIVATEUSE1_NAME: + curr = dist.get_default_backend_for_device(self.device_type) + return f"cpu:gloo,{self.device_type}:{curr}" + return _original_backend.fget(self) + + +common_dtensor.DTensorTestBase.backend = property(_backend) diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py new file mode 100644 index 00000000000..b0fbfadec1f --- /dev/null +++ b/tests/unit_tests/conftest.py @@ -0,0 +1,97 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron bring-up skip configuration for the torchtitan unit test suite. + +This conftest catalogs unit tests that cannot currently run on the Neuron +CPU-eager stack (torch 2.12.x + torch_neuronx 2.12.x) and skips them with a +specific reason so `pytest tests/unit_tests/` finishes green on a Neuron node. +Each skip documents *why* the test is incompatible; the buckets are: + + 1. Gated HuggingFace repos that need auth/network unavailable on the node + (meta-llama, deepseek-ai) -- these fail with 401 GatedRepoError. + 2. spmd_types operations that neuron does not yet support (copy-on-write / + lazy clone), which fail deep inside the collective, not in setup. + 3. loss_parallel-path tests that assert *bitwise* parity between a custom + cross-entropy decomposition and torch's loss_parallel(); the two kernels + diverge at the bit level on the neuron device. + +Optional-dep skips (e.g. `fla`/flash-linear-attention for the qwen3_5 model) +are handled at their own import site with unittest.skipIf, so only the +dependent test is skipped rather than the whole module -- see test_rope.py. + +NOTE: torchvision must match torch. torchvision was upgraded to 0.27.x +alongside torch 2.12 (the stock 0.26 build was compiled for torch 2.11 and +failed to register `torchvision::nms` under 2.12). With 0.27 the flux model +and the multimodal dataset test import and run, so they are no longer skipped. + +The DTensorTestBase distributed tests now run on the neuron device via the +backend shim in this package's __init__.py; only the spmd_types case below +remains skipped. +""" + +import pytest + +# Gated HuggingFace repos: need HF auth/network unavailable on the Neuron node +# (401 GatedRepoError). +_GATED_HF_SKIPS = { + "test_tokenizer.py::TestTokenizerIntegration::test_download_and_build_tokenizer_test_repo_id_meta-llama", + "test_tokenizer.py::TestTokenizerIntegration::test_download_and_build_tokenizer_test_repo_id_deepseek-ai", +} + +# spmd_types operations unsupported on neuron: this test drives +# spmd_types redistribution whose `_replicate_to_varying` calls +# `Tensor._lazy_clone()` (copy-on-write), which returns null storage on the +# neuron device ("Expected storage != nullptr to be true, but got false"). +_SPMD_TYPES_SKIPS = { + "test_parallel_dims.py::TestSpmdLayout::test_partition_spec_order_controls_state_shard", +} + +# loss_parallel path asserts bitwise parity (torch.equal) between a custom +# cross-entropy decomposition and torch's loss_parallel(); the neuron kernels +# are not bitwise-identical to the reference, so these fail on device. Not +# loosened to a tolerance on purpose -- the divergence is real and should be +# quantified before these are enabled. +_LOSS_PARALLEL_SKIPS = { + "test_loss.py::TestLossParallelCrossEntropy::test_loss_parallel_cross_entropy_parity", + "test_loss_dtensor.py::TestCrossEntropyDTensor::test_legacy_tp_loss_parallel", +} + +# Reason strings shared by every entry in each bucket. +_SKIP_BUCKETS = [ + ( + _GATED_HF_SKIPS, + "Gated HuggingFace repo requires auth/network unavailable on the Neuron " + "node (401 GatedRepoError).", + ), + ( + _SPMD_TYPES_SKIPS, + "spmd_types redistribution calls Tensor._lazy_clone() (copy-on-write), " + "unsupported on the neuron device: 'Expected storage != nullptr to be " + "true, but got false'.", + ), + ( + _LOSS_PARALLEL_SKIPS, + "loss_parallel path asserts bitwise parity (torch.equal) with torch's " + "loss_parallel(); the neuron kernels are not bitwise-identical to the " + "reference. Divergence should be quantified before enabling.", + ), +] + + +def pytest_collection_modifyitems(config, items): + """Attach skip markers (with reasons) to known-incompatible unit tests. + + Matching is by nodeid substring so it is stable regardless of the + directory pytest is invoked from (nodeids are prefixed with the + rootdir-relative path, e.g. "tests/unit_tests/test_loss.py::...") and + still matches parametrized tests whose param id is appended to the name + (e.g. the tokenizer repo-id cases). + """ + for item in items: + for nodeids, reason in _SKIP_BUCKETS: + if any(n in item.nodeid for n in nodeids): + item.add_marker(pytest.mark.skip(reason=reason)) + break diff --git a/tests/unit_tests/test_activation_checkpoint.py b/tests/unit_tests/test_activation_checkpoint.py index 360b588c8a4..efc41e7e7ac 100644 --- a/tests/unit_tests/test_activation_checkpoint.py +++ b/tests/unit_tests/test_activation_checkpoint.py @@ -1,15 +1,20 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import unittest from copy import deepcopy import torch from torch.utils.flop_counter import FlopCounterMode -from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC +from torchtitan.distributed.activation_checkpoint import ( + ActivationCheckpointing, + FullAC, + SelectiveAC, +) from torchtitan.models.common.nn_modules import Linear from torchtitan.protocols.module import Module, ModuleDict @@ -266,6 +271,84 @@ def get_recomputed(force_recompute_fqns): # gate and wq still alternate (gate saved, wq recomputed) self.assertEqual(get_recomputed(["output"]), {"wq", "output"}) + def test_matmul_recompute_freq(self): + """matmul_recompute_freq controls which mms are recomputed; 0 saves + every mm. + + mm order in ToyModule: gate (1st), wq (2nd), output (3rd). + """ + from torch.utils._python_dispatch import TorchDispatchMode + + class MmWeightTracker(TorchDispatchMode): + def __init__(self, ptrs): + super().__init__() + self._ptrs = ptrs + self.counts = {n: 0 for n in ptrs.values()} + + def __torch_dispatch__(self, func, types, args, kwargs=None): + if func == torch.ops.aten.mm.default: + for arg in args: + name = self._ptrs.get(arg.data_ptr()) + if name is not None: + self.counts[name] += 1 + break + return func(*args, **(kwargs or {})) + + def get_recomputed(matmul_recompute_freq): + m = ToyModule() + SelectiveAC.Config( + force_recompute_mm_shapes_by_fqns=[], + matmul_recompute_freq=matmul_recompute_freq, + ).build().apply(m) + ptr_to_name = { + mod.weight.data_ptr(): fqn.rsplit(".", 1)[-1] + for fqn, mod in m.named_modules() + if isinstance(mod, Linear) + } + x = torch.randn(64, 512, requires_grad=True) + out = m(x) + tracker = MmWeightTracker(ptr_to_name) + with tracker: + out.backward() + return {n for n, c in tracker.counts.items() if c == 2} + + # Default freq=2 recomputes the 2nd mm (wq): backwards compatible. + self.assertEqual(get_recomputed(2), {"wq"}) + # freq=0 saves every mm (recompute none). + self.assertEqual(get_recomputed(0), set()) + # freq=1 recomputes every mm. + self.assertEqual(get_recomputed(1), {"gate", "wq", "output"}) + # freq=3 recomputes only the 3rd mm (output). + self.assertEqual(get_recomputed(3), {"output"}) + + def test_matmul_recompute_freq_must_be_non_negative(self): + with self.assertRaises(ValueError): + SelectiveAC.Config( + force_recompute_mm_shapes_by_fqns=[], + matmul_recompute_freq=-1, + ).build().apply(ToyModule()) + + def test_selective_ac_policy_captures_no_config(self): + # torch.compile precompile pickles what the guard state reaches, including + # the SAC policy's closure, and a nested Config class does not pickle. + model = ToyModule() + SelectiveAC.Config( + force_recompute_mm_shapes_by_fqns=[], matmul_recompute_freq=3 + ).build().apply(model) + context_fn = model.layers["0"].checkpoint_fn.keywords["context_fn"] + + reached, pending = [], [context_fn] + while pending: + for cell in pending.pop().__closure__ or (): + value = cell.cell_contents + reached.append(value) + if callable(value) and hasattr(value, "__closure__"): + pending.append(value) + self.assertIn(3, reached) + self.assertFalse( + any(isinstance(v, ActivationCheckpointing.Config) for v in reached) + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/test_attn_backend.py b/tests/unit_tests/test_attn_backend.py new file mode 100644 index 00000000000..47c4b8aae3a --- /dev/null +++ b/tests/unit_tests/test_attn_backend.py @@ -0,0 +1,85 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Tests for attention backend selection via ``get_attention_config`` and the +model registries' ``attn_backend`` argument.""" + +import logging +import unittest + +from torchtitan.models.common.attention import ( + FlexAttention, + ScaledDotProductAttention, + VarlenAttention, +) +from torchtitan.models.common.config_utils import get_attention_config + + +class TestGetAttentionConfig(unittest.TestCase): + def test_backends_map_to_config_classes(self): + for backend, cls in ( + ("flex", FlexAttention.Config), + ("varlen", VarlenAttention.Config), + ("sdpa", ScaledDotProductAttention.Config), + ): + self.assertIsInstance(get_attention_config(backend), cls) + + def test_unknown_backend_rejected(self): + with self.assertRaisesRegex(ValueError, "Unknown backend"): + get_attention_config("bogus") + + def test_nki_not_constructed_generically(self): + # nki is gpt_oss-only; its registry builds the wrapper directly. + with self.assertRaisesRegex(ValueError, "Unknown backend"): + get_attention_config("nki") + + def test_sdpa_warns_about_packed_documents(self): + with self.assertLogs(level=logging.WARNING) as captured: + get_attention_config("sdpa") + self.assertTrue( + any("per-document masking" in msg for msg in captured.output), + captured.output, + ) + + +class TestRegistryAttnBackend(unittest.TestCase): + def test_llama3_registry_selects_backend(self): + from torchtitan.models.llama3 import model_registry + + for backend, cls in ( + ("flex", FlexAttention.Config), + ("varlen", VarlenAttention.Config), + ("sdpa", ScaledDotProductAttention.Config), + ): + spec = model_registry("debugmodel", attn_backend=backend) + for layer in spec.model.layers: + self.assertIsInstance(layer.attention.inner_attention, cls) + + def test_gpt_oss_sdpa_rejects_sliding_window(self): + from torchtitan.models.gpt_oss import model_registry + + has_window = any( + layer.attention.sliding_window_size is not None + for layer in model_registry("debugmodel").model.layers + ) + if not has_window: + self.skipTest("debugmodel has no sliding-window layers") + with self.assertRaisesRegex(ValueError, "sliding-window"): + model_registry("debugmodel", attn_backend="sdpa") + + def test_gpt_oss_varlen_window_from_sliding_window_size(self): + from torchtitan.models.gpt_oss import model_registry + + spec = model_registry("debugmodel", attn_backend="varlen") + for layer in spec.model.layers: + sws = layer.attention.sliding_window_size + if sws is not None: + self.assertEqual( + layer.attention.inner_attention.window_size, (sws - 1, 0) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/test_checkpoint.py b/tests/unit_tests/test_checkpoint.py index 1050dbe0032..aa03cdd8207 100644 --- a/tests/unit_tests/test_checkpoint.py +++ b/tests/unit_tests/test_checkpoint.py @@ -14,10 +14,18 @@ from unittest import mock import torch +import torch.distributed.checkpoint as dcp import torch.nn as nn from torch.distributed.checkpoint.state_dict_saver import AsyncSaveResponse from torch.utils.data import DataLoader -from torchtitan.components.checkpoint import CheckpointManager, MODEL, ModelWrapper +from torchtitan.components.checkpoint import ( + CheckpointManager, + FusedQKVLoadPlanner, + MODEL, + ModelWrapper, +) +from torchtitan.models.common.attention import FusedQKVLinear +from torchtitan.models.common.nn_modules import Linear class FakeOptimizersContainer: @@ -324,6 +332,68 @@ def test_load_finds_latest_and_calls_dcp_load(self, mock_load, mock_rank): self.assertTrue(res) manager.close() + @mock.patch("torch.distributed.get_rank", return_value=0) + def test_model_only_load_can_stream_fused_qkv_with_one_dcp_load(self, mock_rank): + source = FusedQKVLinear.Config( + head_dim=2, + n_heads=4, + n_kv_heads=2, + wqkv=Linear.Config(in_features=3, out_features=16, bias=True), + ).build() + destination = FusedQKVLinear.Config( + head_dim=2, + n_heads=4, + n_kv_heads=2, + wqkv=Linear.Config(in_features=3, out_features=16, bias=True), + ).build() + with torch.no_grad(): + source.wqkv.weight.copy_(torch.randn_like(source.wqkv.weight)) + source.wqkv.bias.copy_(torch.randn_like(source.wqkv.bias)) + destination.wqkv.weight.zero_() + destination.wqkv.bias.zero_() + + checkpoint_path = os.path.join(self.test_folder, "seed", "step-0") + dcp.save(source.state_dict(), checkpoint_id=checkpoint_path) + + cfg = self.trainer_config.checkpoint + cfg.folder = "checkpoints" + cfg.initial_load_path = checkpoint_path + cfg.initial_load_model_only = True + cfg.load_only = True + cfg.load_model_state_dict_with_streaming_planner = True + + manager = CheckpointManager( + dataloader=self.data_loader, + model_parts=[destination], + optimizers=self.optimizers, + lr_schedulers=self.lr_schedulers, + states=self.states, + config=cfg, + sd_adapter=None, + base_folder=self.trainer_config.dump_folder, + ) + + with ( + mock.patch( + "torchtitan.components.checkpoint.dcp.load", wraps=dcp.load + ) as mock_load, + mock.patch( + "torchtitan.models.common.attention.torch.cat", + side_effect=AssertionError( + "streaming fused-QKV loading must not allocate a cat buffer" + ), + ), + ): + self.assertTrue(manager.load()) + + mock_load.assert_called_once() + self.assertIsInstance( + mock_load.call_args.kwargs["planner"], FusedQKVLoadPlanner + ) + torch.testing.assert_close(destination.wqkv.weight, source.wqkv.weight) + torch.testing.assert_close(destination.wqkv.bias, source.wqkv.bias) + manager.close() + @mock.patch("torch.distributed.get_rank", return_value=0) @mock.patch("torchtitan.components.checkpoint.dcp.save") @mock.patch("torchtitan.components.checkpoint.dcp.load") @@ -764,10 +834,12 @@ class TestConfigPostInit(unittest.TestCase): def test_valid_default_config(self): """Verify that default values pass initialization.""" try: - CheckpointManager.Config() + cfg = CheckpointManager.Config() except Exception as e: self.fail(f"Default Config raised {type(e).__name__} unexpectedly!") + self.assertFalse(cfg.load_model_state_dict_with_streaming_planner) + def test_sanity_and_range_checks(self): """Test basic field validation like empty strings and negative numbers.""" # Folder cannot be empty @@ -808,6 +880,15 @@ def test_dependency_assertions(self): ) CheckpointManager.Config(initial_load_in_hf=True, initial_load_path=None) + with self.assertRaisesRegex( + ValueError, + "load_model_state_dict_with_streaming_planner requires load_only=True", + ): + CheckpointManager.Config( + load_model_state_dict_with_streaming_planner=True, + load_only=False, + ) + # HF quantized requires HF enabled with self.assertRaisesRegex(ValueError, "requires initial_load_in_hf"): CheckpointManager.Config( @@ -908,6 +989,34 @@ def _split(module, state_dict, prefix, local_metadata): # ... and the in-place refresh picked up the updated parameter. self.assertTrue(torch.all(sd2["a"] == 1.0)) + def test_cache_can_be_disabled_for_load_only(self): + class HookedModule(nn.Module): + def __init__(self): + super().__init__() + self.w = nn.Parameter(torch.zeros(4, 2, 3)) + self.register_state_dict_post_hook(self._split) + + @staticmethod + def _split(module, state_dict, prefix, local_metadata): + w = state_dict.pop(f"{prefix}w") + state_dict[f"{prefix}a"] = w[:, 0].contiguous() + + model = HookedModule() + wrapper = ModelWrapper(model, cache_state_dict=False) + + sd1 = wrapper.state_dict() + ptr_a = sd1["a"].untyped_storage().data_ptr() + self.assertIsNone(wrapper.cached_state_dict) + + with torch.no_grad(): + model.w.fill_(3.0) + + sd2 = wrapper.state_dict() + self.assertIsNot(sd2, sd1) + self.assertNotEqual(sd2["a"].untyped_storage().data_ptr(), ptr_a) + self.assertTrue(torch.all(sd2["a"] == 3.0)) + self.assertIsNone(wrapper.cached_state_dict) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/test_compiled_grad_clip.py b/tests/unit_tests/test_compiled_grad_clip.py new file mode 100644 index 00000000000..1f499bc17b6 --- /dev/null +++ b/tests/unit_tests/test_compiled_grad_clip.py @@ -0,0 +1,131 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +from unittest import mock + +import torch + +from torchtitan.config import CompileConfig +from torchtitan.distributed import utils as dist_utils +from torchtitan.tools import logging as logging_utils + + +def test_compiled_grad_clip_uses_configured_backend(monkeypatch): + compile_calls = [] + + def compile_fn(fn, **kwargs): + compile_calls.append((fn, kwargs)) + return fn + + monkeypatch.setattr(torch, "compile", compile_fn) + monkeypatch.setattr(dist_utils.dist, "all_reduce", mock.Mock()) + + parameter = torch.nn.Parameter(torch.tensor([3.0, 4.0])) + parameter.grad = torch.tensor([3.0, 4.0]) + clipper = dist_utils.CompiledGradClip(backend="aot_eager") + + total_norm = dist_utils.clip_grad_norm_( + [parameter], + max_norm=1.0, + compiled_grad_clip=clipper, + ) + + torch.testing.assert_close(total_norm, torch.tensor(5.0)) + torch.testing.assert_close(parameter.grad, torch.tensor([0.6, 0.8])) + assert compile_calls == [ + ( + dist_utils._compute_grouped_norm_sq, + {"backend": "aot_eager", "dynamic": False}, + ), + ( + dist_utils._compute_clip_coef, + {"backend": "aot_eager", "dynamic": False}, + ), + ] + + +def test_compiled_grad_clip_is_opt_in(): + parameter = torch.nn.Parameter(torch.tensor([3.0, 4.0])) + parameter.grad = torch.tensor([3.0, 4.0]) + compiled_clip = mock.Mock(return_value=torch.tensor(7.0)) + + eager_total_norm = dist_utils.clip_grad_norm_([parameter], max_norm=1.0) + + torch.testing.assert_close(eager_total_norm, torch.tensor(5.0)) + compiled_clip.assert_not_called() + + parameter.grad = torch.tensor([3.0, 4.0]) + compiled_total_norm = dist_utils.clip_grad_norm_( + [parameter], + max_norm=1.0, + compiled_grad_clip=compiled_clip, + ) + + torch.testing.assert_close(compiled_total_norm, torch.tensor(7.0)) + compiled_clip.assert_called_once_with([parameter], 1.0) + + +def test_build_compiled_grad_clip_is_opt_in(): + assert ( + dist_utils.build_compiled_grad_clip( + CompileConfig(enable=True, backend="aot_eager") + ) + is None + ) + + clipper = dist_utils.build_compiled_grad_clip( + CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="aot_eager", + ) + ) + + assert isinstance(clipper, dist_utils.CompiledGradClip) + assert clipper._backend == "aot_eager" + + +def test_compiled_grad_clip_warns_once_when_falling_back( + monkeypatch, + caplog, +): + monkeypatch.setattr(logging_utils, "_logged", set()) + parameter = torch.nn.Parameter(torch.tensor([3.0, 4.0])) + compiled_clip = mock.Mock() + + with caplog.at_level("WARNING"): + for _ in range(2): + parameter.grad = torch.tensor([3.0, 4.0]) + total_norm = dist_utils.clip_grad_norm_( + [parameter], + max_norm=1.0, + norm_type=1.0, + compiled_grad_clip=compiled_clip, + ) + torch.testing.assert_close(total_norm, torch.tensor(7.0)) + + compiled_clip.assert_not_called() + warnings = [ + record.message + for record in caplog.records + if "Compiled gradient clipping was requested" in record.message + ] + assert len(warnings) == 1 + assert "norm_type is 1.0, not 2.0" in warnings[0] + assert "falling back to eager clipping" in warnings[0] + + +def test_compiled_grad_clip_without_local_grads_still_reduces(monkeypatch): + all_reduce = mock.Mock(side_effect=lambda tensor, **_: tensor.fill_(25.0)) + monkeypatch.setattr(dist_utils.dist, "all_reduce", all_reduce) + + parameter = torch.nn.Parameter(torch.tensor([1.0])) + clipper = dist_utils.CompiledGradClip(backend="aot_eager") + + total_norm = clipper([parameter], max_norm=1.0) + + torch.testing.assert_close(total_norm, torch.tensor(5.0)) + assert total_norm.device == parameter.device + all_reduce.assert_called_once() diff --git a/tests/unit_tests/test_context_parallel_rope.py b/tests/unit_tests/test_context_parallel_rope.py new file mode 100644 index 00000000000..0f655bb4903 --- /dev/null +++ b/tests/unit_tests/test_context_parallel_rope.py @@ -0,0 +1,217 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +import tempfile +import unittest +from unittest.mock import Mock, patch + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch.distributed.tensor import DTensor, Replicate, Shard + +from torchtitan.distributed.context_parallel import ( + apply_cp_to_rope, + prepare_context_parallel_input, +) +from torchtitan.models.common.rope import ComplexRoPE, CosSinRoPE + + +def _cp_mesh(rank, size): + mesh = Mock(spec=DeviceMesh) + mesh.ndim = 1 + mesh.mesh_dim_names = ("cp",) + mesh.size.return_value = size + mesh.get_local_rank.return_value = rank + return mesh + + +class TestContextParallelRoPE(unittest.TestCase): + def test_input_positions_are_resolved_before_sharding(self): + inputs = torch.arange(16).view(2, 8) + labels = inputs + 1 + packed_positions = torch.tensor([[0, 1, 2, 0, 1, 2, 3, 0]]).expand(2, -1) + + def shard_second_half(mesh, buffers, masks, load_balancer): + return tuple(buffer[:, 4:] for buffer in buffers), masks + + for kwargs in ({}, {"positions": None}, {"positions": packed_positions}): + with self.subTest(positions=kwargs.get("positions")): + expected = kwargs.get("positions") + if expected is None: + expected = torch.arange(8).expand(2, -1) + with patch( + "torchtitan.distributed.context_parallel.cp_shard", + side_effect=shard_second_half, + ): + ( + local_inputs, + local_labels, + local_kwargs, + ) = prepare_context_parallel_input( + inputs, + labels, + kwargs, + _cp_mesh(1, 2), + torch.device("cpu"), + None, + ) + torch.testing.assert_close(local_inputs, inputs[:, 4:]) + torch.testing.assert_close(local_labels, labels[:, 4:]) + torch.testing.assert_close(local_kwargs["positions"], expected[:, 4:]) + + def test_shards_match_global_rotation_and_gradients(self): + """Every rank must rotate at global positions, including after reinit.""" + torch.manual_seed(42) + for rope_cls in ( + ComplexRoPE, + CosSinRoPE, + ): + for cp_size in (1, 2, 16): + for load_balancer in (None, "headtail"): + with self.subTest( + rope=rope_cls.__name__, cp=cp_size, balancer=load_balancer + ): + seq_len = 64 + config = rope_cls.Config( + dim=8, + max_seq_len=128, + use_prefix_cache_for_positions=True, + ) + reference = rope_cls(config) + query = torch.randn(2, seq_len, 4, 8, requires_grad=True) + key = torch.randn(2, seq_len, 2, 8, requires_grad=True) + query_ref, key_ref = reference(query, key) + (query_ref.sum() + key_ref.sum()).backward() + for rank in range(cp_size): + indices = torch.arange(seq_len) + if cp_size > 1 and load_balancer == "headtail": + chunks = indices.chunk(2 * cp_size) + indices = torch.cat( + (chunks[rank], chunks[2 * cp_size - rank - 1]) + ) + else: + indices = indices.chunk(cp_size)[rank] + rope = rope_cls(config) + apply_cp_to_rope( + torch.nn.Sequential(rope), + _cp_mesh(rank, cp_size), + load_balancer_type=load_balancer, + ) + rope.to_empty(device="cpu") + rope.init_states() + query_local = query.detach()[:, indices].requires_grad_() + key_local = key.detach()[:, indices].requires_grad_() + # Prefix-cache semantics intentionally ignore document + # resets, just as the CP1 baseline does. + positions = torch.zeros(2, len(indices), dtype=torch.long) + query_out, key_out = rope(query_local, key_local, positions) + torch.testing.assert_close( + query_out, query_ref[:, indices], rtol=0, atol=0 + ) + torch.testing.assert_close( + key_out, key_ref[:, indices], rtol=0, atol=0 + ) + (query_out.sum() + key_out.sum()).backward() + torch.testing.assert_close( + query_local.grad, query.grad[:, indices], rtol=0, atol=0 + ) + torch.testing.assert_close( + key_local.grad, key.grad[:, indices], rtol=0, atol=0 + ) + + def test_position_gather_is_unchanged(self): + rope = CosSinRoPE(CosSinRoPE.Config(dim=8, max_seq_len=64)) + query = torch.randn(2, 4, 2, 8) + positions = torch.tensor([[10, 11, 0, 1], [2, 3, 4, 0]]) + expected = rope(query, query, positions) + apply_cp_to_rope(rope, _cp_mesh(1, 2), load_balancer_type="ptrr") + actual = rope(query, query, positions) + for result, reference in zip(actual, expected): + torch.testing.assert_close(result, reference, rtol=0, atol=0) + + def test_prefix_cache_rejects_ptrr_during_setup(self): + for rope_cls, prefix_config in ( + (ComplexRoPE, True), + (CosSinRoPE, True), + ): + with self.subTest(rope=rope_cls.__name__, prefix_config=prefix_config): + rope = rope_cls( + rope_cls.Config( + dim=8, + max_seq_len=64, + use_prefix_cache_for_positions=prefix_config, + ) + ) + with self.assertRaisesRegex(ValueError, "Prefix RoPE caches require"): + apply_cp_to_rope(rope, _cp_mesh(1, 2), load_balancer_type="ptrr") + + def test_mesh_axis_is_validated(self): + rope = ComplexRoPE(ComplexRoPE.Config(dim=8, max_seq_len=64)) + mesh = _cp_mesh(0, 2) + mesh.mesh_dim_names = ("tp",) + with self.assertRaisesRegex(ValueError, "mesh named 'cp'"): + apply_cp_to_rope(rope, mesh, load_balancer_type=None) + mesh.ndim = 2 + mesh.mesh_dim_names = ("dp", "cp") + with self.assertRaisesRegex(ValueError, "1D mesh"): + apply_cp_to_rope(rope, mesh, load_balancer_type=None) + + +class TestContextParallelRoPEDTensor(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls._temp_dir = tempfile.TemporaryDirectory() + dist.init_process_group( + backend="gloo", + init_method=f"file://{cls._temp_dir.name}/store", + world_size=1, + rank=0, + ) + + @classmethod + def tearDownClass(cls): + dist.destroy_process_group() + cls._temp_dir.cleanup() + + def test_tp_and_multi_axis_cache_placements(self): + """Cache broadcasts over batch/heads and follows CP sequence sharding.""" + for rope_cls in (ComplexRoPE, CosSinRoPE): + for axis_names, placements in ( + (("tp",), (Shard(2),)), + (("dp", "cp", "tp"), (Shard(0), Shard(1), Shard(2))), + ): + with self.subTest(rope=rope_cls.__name__, axes=axis_names): + mesh = init_device_mesh( + "cpu", (1,) * len(axis_names), mesh_dim_names=axis_names + ) + rope = rope_cls( + rope_cls.Config( + dim=8, max_seq_len=64, use_prefix_cache_for_positions=True + ) + ) + expected = rope.cache[4:8].view(1, 4, 1, -1).clone() + rope.cache = DTensor.from_local( + rope.cache, mesh, (Replicate(),) * mesh.ndim, run_check=False + ) + apply_cp_to_rope(rope, _cp_mesh(1, 2), load_balancer_type=None) + query = DTensor.from_local( + torch.randn(2, 4, 2, 8), mesh, placements, run_check=False + ) + cache = rope._reshape_cache(query) + torch.testing.assert_close( + cache.to_local(), expected, rtol=0, atol=0 + ) + self.assertEqual( + cache.placements, + tuple( + Shard(1) if axis == "cp" else Replicate() + for axis in axis_names + ), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/test_loss.py b/tests/unit_tests/test_loss.py index 6165d475a44..48b5e10b910 100644 --- a/tests/unit_tests/test_loss.py +++ b/tests/unit_tests/test_loss.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import unittest @@ -313,31 +314,29 @@ def test_loss_parallel_cross_entropy_parity(self): ): # Build global test data once, then let DTensor derive the # exact local shards for TP-only and DP+TP placements. - generator = torch.Generator(device=self.device_type).manual_seed( - seed - ) + # Seed a CPU generator and move to the test device: some + # accelerators (e.g. neuron) do not provide a device-side + # RNG generator. Determinism is preserved by the fixed seed. + generator = torch.Generator().manual_seed(seed) global_logits = torch.randn( B, L, vocab_size, - device=self.device_type, generator=generator, - ) + ).to(self.device_type) global_labels = torch.randint( 0, vocab_size, (B, L), - device=self.device_type, dtype=torch.long, generator=generator, - ) + ).to(self.device_type) if ignore: mask = torch.rand( B, L, - device=self.device_type, generator=generator, - ) + ).to(self.device_type) global_labels[mask < 0.3] = IGNORE_INDEX # distribute as DTensor; loss_parallel wrapper takes DTensors. logits_dtensor = distribute_tensor( diff --git a/tests/unit_tests/test_loss_dtensor.py b/tests/unit_tests/test_loss_dtensor.py index b63ffebaa6a..6a3203c18c8 100644 --- a/tests/unit_tests/test_loss_dtensor.py +++ b/tests/unit_tests/test_loss_dtensor.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import contextlib import unittest @@ -48,10 +49,14 @@ def _run_one(self, *, full_dtensor: bool, tp_shard_v: bool): labels_pl = (Replicate(),) B, S, V = 4, 16, 64 - gen = torch.Generator(device=self.device_type).manual_seed(42) - global_pred = torch.randn(B, S, V, device=self.device_type, generator=gen) - global_labels = torch.randint( - 0, V, (B, S), device=self.device_type, dtype=torch.long, generator=gen + # Seed a CPU generator and move to the test device: some accelerators + # (e.g. neuron) do not provide a device-side RNG generator. The parity + # checks below compare two computations over the same data, so the + # exact RNG stream does not matter as long as it is deterministic. + gen = torch.Generator().manual_seed(42) + global_pred = torch.randn(B, S, V, generator=gen).to(self.device_type) + global_labels = torch.randint(0, V, (B, S), dtype=torch.long, generator=gen).to( + self.device_type ) ref_loss = _reference_ce(global_pred, global_labels) diff --git a/tests/unit_tests/test_moe_global_lbl.py b/tests/unit_tests/test_moe_global_lbl.py new file mode 100644 index 00000000000..cb9136f7fdf --- /dev/null +++ b/tests/unit_tests/test_moe_global_lbl.py @@ -0,0 +1,326 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""The global-batch load-balancing loss (moe.attach_global_batch_lbl), on CPU.""" + +import unittest + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor, Partial +from torch.testing._internal.common_utils import run_tests +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) + +from torchtitan.models.common.config_utils import ( + make_experts_config, + make_moe_config, + make_router_config, +) +from torchtitan.models.common.moe import configure_global_batch_lbl + + +DIM, NUM_EXPERTS, TOP_K, COEFF = 16, 8, 2, 0.5 +_HAS_FX_TRACER = hasattr(torch.compiler, "_patch_engine_backward") + + +def _moe(load_balancing="global_aux_loss"): + init = {"weight": lambda x: nn.init.normal_(x, std=0.5)} + config = make_moe_config( + num_experts=NUM_EXPERTS, + router=make_router_config( + dim=DIM, + num_experts=NUM_EXPERTS, + top_k=TOP_K, + gate_param_init=init, + score_func="softmax", + ), + experts=make_experts_config( + dim=DIM, + hidden_dim=2 * DIM, + num_experts=NUM_EXPERTS, + top_k=TOP_K, + param_init=init, + comm_backend="standard", + ), + load_balance_coeff=COEFF, + ) + config.load_balancing = load_balancing + torch.manual_seed(0) + moe = config.build() + with torch.no_grad(): + for param in moe.parameters(): + nn.init.normal_(param, std=0.5) + return moe.train() + + +def _reference_fp64(moe, x_BLD, num_microbatches=1): + """coeff * E * sum_i f_i P_i from the router alone, in fp64, as a function of the gate.""" + gate = moe.router.gate.weight.detach().double().requires_grad_() + probs_TE = torch.softmax(x_BLD.reshape(-1, DIM).double() @ gate.t(), dim=-1) + top_TK = probs_TE.topk(TOP_K, dim=-1).indices + counts_E = torch.zeros(NUM_EXPERTS, dtype=torch.float64).index_add_( + 0, top_TK.flatten(), torch.ones(top_TK.numel(), dtype=torch.float64) + ) + f_E = counts_E / counts_E.sum() + aux = COEFF * NUM_EXPERTS * (f_E * probs_TE.mean(dim=0)).sum() / num_microbatches + aux.backward() + return gate.grad + + +def _aux_gate_grad(moe, x_BLD): + """Gate gradient of the attached loss alone (a zero main loss).""" + (moe(x_BLD) * 0).sum().backward() + return moe.router.gate.weight.grad + + +def test_matches_the_paper_formula(): + x_BLD = torch.randn(2, 8, DIM) + moe = _moe() + torch.testing.assert_close( + _aux_gate_grad(moe, x_BLD).double(), + _reference_fp64(moe, x_BLD), + rtol=1e-5, + atol=1e-7, + ) + + +def test_micro_batches_share_the_step_normalization(): + x_BLD = torch.randn(2, 8, DIM) + moe = _moe() + moe.lbl_num_microbatches = 2 + torch.testing.assert_close( + _aux_gate_grad(moe, x_BLD).double(), + _reference_fp64(moe, x_BLD, num_microbatches=2), + rtol=1e-5, + atol=1e-7, + ) + + +def test_expert_bias_default_unchanged(): + moe = _moe("expert_bias") + assert moe.expert_bias_E is not None + grad = _aux_gate_grad(moe, torch.randn(2, 8, DIM)) + assert grad is None or float(grad.abs().max()) == 0.0 + + +def test_activation_checkpointing_unchanged(): + from torch.utils.checkpoint import checkpoint + + x_BLD = torch.randn(2, 8, DIM) + moe = _moe() + want = _aux_gate_grad(moe, x_BLD).clone() + for wrap in ( + lambda m, x: checkpoint(m, x, use_reentrant=False), + lambda m, x: torch.compile( + lambda y: checkpoint(m, y, use_reentrant=False), backend="aot_eager" + )(x), + ): + moe = _moe() + (wrap(moe, x_BLD) * 0).sum().backward() + torch.testing.assert_close(moe.router.gate.weight.grad, want) + + +def test_unknown_mechanism_refused(): + with pytest.raises(ValueError, match="unknown load_balancing"): + _moe("global_aux") + + +class TestTwoRanksEqualOne(DTensorTestBase): + @property + def world_size(self) -> int: + return 2 + + @property + def device_type(self) -> str: + return "cpu" + + @with_comms + def test_router_grad(self): + torch.manual_seed(1) + x_BLD = torch.randn(4, 8, DIM) + single = _moe() + want = _aux_gate_grad(single, x_BLD) + + moe = _moe() + moe.lbl_batch_mesh = init_device_mesh("cpu", (self.world_size,)) + grad = _aux_gate_grad(moe, x_BLD.chunk(self.world_size)[self.rank]).clone() + # Summed over ranks, as the main loss's gradients are. + dist.all_reduce(grad) + assert float(grad.norm()) > 0 + torch.testing.assert_close(grad, want) + + @unittest.skipIf( + not _HAS_FX_TRACER, "graph_trainer's minimal_fx_tracer needs torch >= 2.13" + ) + @with_comms + def test_router_grad_traced_by_aot_fx_trace(self): + # The counts all-reduce runs inside the make_fx graph (functional collective). + from torchtitan.experiments.graph_trainer.make_fx_tracer import ( + minimal_fx_tracer, + run_traced, + ) + + torch.manual_seed(1) + x_BLD = torch.randn(4, 8, DIM) + want = _aux_gate_grad(_moe(), x_BLD) + moe = _moe() + moe.lbl_batch_mesh = init_device_mesh("cpu", (self.world_size,)) + + def fwd_bwd(x): + loss = (moe(x) * 0).sum() + return [loss, *torch.autograd.grad(loss, [moe.router.gate.weight])] + + local = x_BLD.chunk(self.world_size)[self.rank] + traced = minimal_fx_tracer(fwd_bwd, module=moe)(local) + grad = run_traced(traced, module=moe)(local)[1].clone() + dist.all_reduce(grad) + torch.testing.assert_close(grad, want) + + @with_comms + def test_with_a_dtensor_main_loss(self): + # Under SimpleFSDP the CE loss is a DTensor (Partial over the batch axis). + mesh = init_device_mesh("cpu", (self.world_size,)) + torch.manual_seed(1) + x_BLD = torch.randn(4, 8, DIM) + tokens = x_BLD.shape[0] * x_BLD.shape[1] + single = _moe() + (single(x_BLD).sum() / tokens).backward() + + moe = _moe() + moe.lbl_batch_mesh = mesh + out = moe(x_BLD.chunk(self.world_size)[self.rank]) + DTensor.from_local(out.sum() / tokens, mesh, [Partial()]).backward() + grad = moe.router.gate.weight.grad.clone() + dist.all_reduce(grad) + torch.testing.assert_close(grad, single.router.gate.weight.grad) + + +def test_traced_by_aot_fx_trace(): + """GraphTrainer's aot_fx_trace traces fwd+bwd itself; the attached term must be in that graph.""" + if not _HAS_FX_TRACER: + pytest.skip("graph_trainer's minimal_fx_tracer needs torch >= 2.13") + from torchtitan.experiments.graph_trainer.make_fx_tracer import ( + minimal_fx_tracer, + run_traced, + ) + + x_BLD = torch.randn(2, 8, DIM) + want = _aux_gate_grad(_moe(), x_BLD) + moe = _moe() + + def fwd_bwd(x): + loss = (moe(x) * 0).sum() + return [loss, *torch.autograd.grad(loss, [moe.router.gate.weight])] + + traced = minimal_fx_tracer(fwd_bwd, module=moe)(x_BLD) + gate_grad = run_traced(traced, module=moe)(x_BLD)[1] + assert float(want.norm()) > 0 + torch.testing.assert_close(gate_grad, want) + + +class _Stage(nn.Module): + """One pipeline stage: a residual global-batch-LBL MoE block, as in a decoder layer.""" + + def __init__(self, seed): + super().__init__() + torch.manual_seed(seed) + self.moe = _moe() + + def forward(self, x_BLD): + return x_BLD + self.moe(x_BLD) + + +def _pp_loss(main_weight, global_tokens): + # Scaled by the step's token count, as the trainer's loss is (scale_grads=False). + def loss_fn(out_BLD, target_BLD): + return main_weight * (out_BLD * target_BLD).sum() / global_tokens + + return loss_fn + + +class TestPipelineParallel(DTensorTestBase): + """Two PP stages equal one process running the same micro-batches back to back.""" + + @property + def world_size(self) -> int: + return 2 + + @property + def device_type(self) -> str: + return "cpu" + + @with_comms + def test_router_grads_match_the_unpipelined_step(self): + from torch.distributed.pipelining import ( + PipelineStage, + Schedule1F1B, + ScheduleGPipe, + ScheduleInterleaved1F1B, + ScheduleInterleavedZeroBubble, + ) + + torch.manual_seed(3) + x_BLD, target_BLD = torch.randn(4, 8, DIM), torch.randn(4, 8, DIM) + tokens = x_BLD.shape[0] * x_BLD.shape[1] + # (schedule, stages per rank); the zero-bubble schedule splits input and weight backward. + for schedule_cls, n_local in ( + (ScheduleGPipe, 1), + (Schedule1F1B, 1), + (ScheduleInterleaved1F1B, 2), + (ScheduleInterleavedZeroBubble, 2), + ): + n_stages = self.world_size * n_local + stage_ids = [self.rank + i * self.world_size for i in range(n_local)] + for n_microbatches in (2, 4): + for main_weight in (0.0, 1.0): + loss_fn = _pp_loss(main_weight, tokens) + # Reference: every stage in one process, one forward/backward per micro-batch. + ref = [_Stage(i) for i in range(n_stages)] + configure_global_batch_lbl( + ref, batch_mesh=None, num_forwards=n_microbatches + ) + for x, t in zip( + x_BLD.chunk(n_microbatches), target_BLD.chunk(n_microbatches) + ): + for s in ref: + x = s(x) + loss_fn(x, t).backward() + + stages = [_Stage(i) for i in stage_ids] + configure_global_batch_lbl( + stages, batch_mesh=None, num_forwards=n_microbatches + ) + pipe = [ + PipelineStage(s, i, n_stages, torch.device("cpu")) + for s, i in zip(stages, stage_ids) + ] + schedule = schedule_cls( + pipe if n_local > 1 else pipe[0], + n_microbatches=n_microbatches, + loss_fn=loss_fn, + scale_grads=False, + ) + if self.rank == 0: + schedule.step(x_BLD) + else: + schedule.step(target=target_BLD, losses=[]) + for s, i in zip(stages, stage_ids): + want = ref[i].moe.router.gate.weight.grad + assert float(want.norm()) > 0 + torch.testing.assert_close( + s.moe.router.gate.weight.grad, + want, + msg=schedule_cls.__name__, + ) + + +if __name__ == "__main__": + run_tests() diff --git a/tests/unit_tests/test_moe_tp_input_grads.py b/tests/unit_tests/test_moe_tp_input_grads.py new file mode 100644 index 00000000000..30c2dd3fdf3 --- /dev/null +++ b/tests/unit_tests/test_moe_tp_input_grads.py @@ -0,0 +1,76 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""With EP off, the routed experts' input gradients must be declared Partial over TP. + +Each TP rank computes its experts on a hidden-dim slice, so its x / top-k-score gradients are partial +sums; without the declaration local_map assumes Replicate and never sums them. +""" + +from __future__ import annotations + +import pytest +import spmd_types as spmd + +from torchtitan.models.common.moe import MoE +from torchtitan.models.common.moe_sharding import set_moe_sharding_config, TP + +_EXPERT_LAYOUT = {"w1_EFD": spmd.S(1), "w2_EDF": spmd.S(2), "w3_EFD": spmd.S(1)} + + +def _moe_config() -> MoE.Config: + from torchtitan.models.common.config_utils import ( + make_experts_config, + make_moe_config, + make_router_config, + ) + + init = {"weight": lambda x: x} + return make_moe_config( + num_experts=8, + router=make_router_config( + dim=16, + num_experts=8, + top_k=2, + gate_param_init=init, + score_func="softmax", + ), + experts=make_experts_config( + dim=16, + hidden_dim=32, + num_experts=8, + top_k=2, + param_init=init, + comm_backend="standard", + ), + ) + + +@pytest.mark.parametrize("enable_sp", [True, False]) +def test_tp_only_experts_sum_their_input_gradients_over_tp(enable_sp): + """KILLS: the TP-only MoE handing back ~1/TP of its input and routing-score gradients.""" + cfg = _moe_config() + set_moe_sharding_config( + cfg, enable_ep=False, enable_sp=enable_sp, expert_param_layout=_EXPERT_LAYOUT + ) + grads = cfg.experts.sharding_config.local_map.in_grad_placements + assert ( + grads is not None + ), "TP-only experts omit in_grad_placements: partials are never summed" + x_grad, scores_grad = grads[0], grads[1] + assert x_grad.axis_types[TP] == spmd.P + assert scores_grad.axis_types[TP] == spmd.P + + +def test_the_ep_path_is_unchanged(): + """The EP path already declared its gradients; the fix must not change them.""" + cfg = _moe_config() + set_moe_sharding_config( + cfg, enable_ep=True, enable_sp=True, expert_param_layout=_EXPERT_LAYOUT + ) + grads = cfg.experts.sharding_config.local_map.in_grad_placements + assert grads is not None + assert grads[0] == grads[1] == grads[2] + assert grads[0].axis_types[TP] != spmd.P diff --git a/tests/unit_tests/test_moe_tp_only_input_grads_numeric.py b/tests/unit_tests/test_moe_tp_only_input_grads_numeric.py new file mode 100644 index 00000000000..47aff9591d4 --- /dev/null +++ b/tests/unit_tests/test_moe_tp_only_input_grads_numeric.py @@ -0,0 +1,226 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""TP-only MoE input and router gradients at TP=N must match TP=1 (not ~1/N of it). + +Real MoE through set_moe_sharding_config + Module.parallelize on a CPU/gloo TP mesh. Tolerance is bf16: +the experts run torch._grouped_mm in bf16, so each rank's partial sum is rounded before the reduce. +""" + +import copy +from functools import partial +from unittest.mock import patch + +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed.tensor import DTensor, Shard +from torch.testing._internal.common_utils import run_tests +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) + +from torchtitan.distributed.parallel_dims import ParallelDims +from torchtitan.models.common.config_utils import ( + make_experts_config, + make_moe_config, + make_router_config, +) +from torchtitan.models.common.moe import MoE +from torchtitan.models.common.moe_sharding import set_moe_sharding_config +from torchtitan.models.qwen3.sharding import _GROUPED_EXPERTS_PARAM_LAYOUT +from torchtitan.protocols.sharding import LocalMapConfig + + +_DIM = 16 +# CPU torch._grouped_mm needs 16-byte strides, i.e. bf16 sizes that are multiples of 8. Each rank +# holds _HIDDEN_DIM / TP of the hidden dim, so 32 covers TP2 and TP4. +_HIDDEN_DIM = 32 +_NUM_EXPERTS = 8 +_TOP_K = 2 +# Sequence parallelism needs the sequence to split evenly over TP. +_BATCH, _SEQ_LEN = 2, 8 +# Bound on max|TP - TP1| / max|TP1|; bf16-limited, see the module docstring. +_TOLERANCE = 1e-2 + + +def _moe_config(*, enable_sp: bool) -> MoE.Config: + """A small Qwen3-style MoE (softmax router with route_norm, top-2), sharded as Qwen3 shards it.""" + cfg = make_moe_config( + num_experts=_NUM_EXPERTS, + router=make_router_config( + dim=_DIM, + num_experts=_NUM_EXPERTS, + top_k=_TOP_K, + gate_param_init={"weight": partial(nn.init.normal_, std=_DIM**-0.5)}, + score_func="softmax", + route_norm=True, + ), + experts=make_experts_config( + dim=_DIM, + hidden_dim=_HIDDEN_DIM, + num_experts=_NUM_EXPERTS, + top_k=_TOP_K, + param_init={ + "w1_EFD": partial(nn.init.normal_, std=_DIM**-0.5), + "w2_EDF": partial(nn.init.normal_, std=_HIDDEN_DIM**-0.5), + "w3_EFD": partial(nn.init.normal_, std=_DIM**-0.5), + }, + # Falls back to local dispatch when EP is off. + comm_backend="standard", + ), + ) + set_moe_sharding_config( + cfg, + enable_ep=False, + enable_sp=enable_sp, + expert_param_layout=_GROUPED_EXPERTS_PARAM_LAYOUT, + ) + return cfg + + +def _error_and_share(got: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """(max|got - ref| / max|ref|, / ): how far off, and how much of ref it holds.""" + error = (got - ref).abs().max() / ref.abs().max() + share = (got * ref).sum() / (ref * ref).sum() + return error.item(), share.item() + + +class TestTPOnlyMoEInputGrads(DTensorTestBase): + @property + def device_type(self) -> str: + # The bound was measured on CPU, and what is under test is DTensor gradient bookkeeping, not + # accelerator kernels. + return "cpu" + + @property + def world_size(self) -> int: + # The smallest mesh with a TP partial to lose. + return 2 + + def _parallel_dims(self) -> ParallelDims: + # TP only: EP, DP, CP and PP are all off. + parallel_dims = ParallelDims( + dp_replicate=1, + dp_shard=1, + cp=1, + tp=self.world_size, + pp=1, + ep=1, + world_size=self.world_size, + ) + with patch( + "torchtitan.distributed.parallel_dims.device_type", self.device_type + ): + parallel_dims.build_mesh() + return parallel_dims + + def _gradients( + self, cfg: MoE.Config, *, enable_sp: bool + ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """TP1 and TP gradients of the MoE input and router gate weight, for one input and upstream grad. + + Under SP the TP module takes and returns this rank's sequence shard, so its input gradient is + that shard and the TP1 one is sliced to match. Gate gradients are full tensors. + """ + gen = torch.Generator().manual_seed(0) + x_BLD = torch.randn(_BATCH, _SEQ_LEN, _DIM, generator=gen) + grad_out_BLD = torch.randn(_BATCH, _SEQ_LEN, _DIM, generator=gen) + + torch.manual_seed(0) + reference = cfg.build() + reference.init_states() + sharded = copy.deepcopy(reference) + + x_ref = x_BLD.clone().requires_grad_() + reference(x_ref).backward(grad_out_BLD) + + parallel_dims = self._parallel_dims() + sharded.parallelize(parallel_dims) + # The precondition for TP partials: each rank holds a hidden-dim slice of every expert. + self.assertEqual( + ( + sharded.experts.w1_EFD.placements, + sharded.experts.w2_EDF.placements, + sharded.experts.w3_EFD.placements, + ), + ((Shard(1),), (Shard(2),), (Shard(1),)), + ) + tp_rank = parallel_dims.get_mesh("tp").get_local_rank() + + def local(t: torch.Tensor) -> torch.Tensor: + return t.chunk(self.world_size, dim=1)[tp_rank] if enable_sp else t + + x_tp = local(x_BLD).clone().requires_grad_() + out = sharded(x_tp) + self.assertIsInstance(out, DTensor) + out.to_local().backward(local(grad_out_BLD)) + + gate_grad = sharded.router.gate.weight.grad + self.assertIsInstance(gate_grad, DTensor) + reference_grads = { + "input": local(x_ref.grad), + "router gate weight": reference.router.gate.weight.grad, + } + tp_grads = {"input": x_tp.grad, "router gate weight": gate_grad.full_tensor()} + return reference_grads, tp_grads + + @with_comms + def test_input_and_router_gradients_match_tp1(self): + """KILLS: the TP-only MoE handing back ~1/TP of its input and router gradients. + + The declaration from before the fix, ``in_grad_placements=None``, runs last as the negative + control: each rank's gradients are then far from TP1, yet they sum over TP to it, i.e. the + partials were computed and never reduced. Measured share of the TP1 input gradient, + / : 0.59/0.41 at TP2 and 0.27/0.32/0.21/0.20 at TP4. + """ + for enable_sp in (False, True): + reference, sharded = self._gradients( + _moe_config(enable_sp=enable_sp), enable_sp=enable_sp + ) + for name, ref in reference.items(): + error, share = _error_and_share(sharded[name], ref) + self.assertLess( + error, + _TOLERANCE, + f"{name} gradient at TP{self.world_size} (enable_sp={enable_sp}) is off by " + f"{error:.3g} of its magnitude and holds {share:.3f} of the TP1 gradient; " + "~1/TP means the routed experts' TP partials were never summed", + ) + + # Negative control, without SP so that each rank's gradient is exactly its own partial. + cfg = _moe_config(enable_sp=False) + cfg.experts.sharding_config.local_map = LocalMapConfig(in_grad_placements=None) + reference, sharded = self._gradients(cfg, enable_sp=False) + for name, ref in reference.items(): + error, share = _error_and_share(sharded[name], ref) + # Ten times the tolerance: a missing partial, not rounding. + self.assertGreater( + error, + 10 * _TOLERANCE, + f"{name} gradient matches TP1 (share {share:.3f}) without in_grad_placements; if " + "local_map now infers Partial, this test no longer proves the fix is needed", + ) + # TP only, so the default process group is the TP group. + total = sharded[name].clone() + dist.all_reduce(total) + total_error, _ = _error_and_share(total, ref) + self.assertLess( + total_error, + _TOLERANCE, + f"{name} gradients without in_grad_placements do not sum over TP to TP1", + ) + + +class TestTPOnlyMoEInputGradsAtFour(TestTPOnlyMoEInputGrads): + @property + def world_size(self) -> int: + # The TP degree of the GPU runs where every ffn_norm gradient came out 0.26-0.28x its TP1 value. + return 4 + + +if __name__ == "__main__": + run_tests() diff --git a/tests/unit_tests/test_profiler.py b/tests/unit_tests/test_profiler.py index 9c4208cfd7b..629ef970d10 100644 --- a/tests/unit_tests/test_profiler.py +++ b/tests/unit_tests/test_profiler.py @@ -1,12 +1,20 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. +import enum +import os +import sys +import tempfile +import types import unittest from unittest import mock +import torch + from torchtitan.tools.profiler import Profiler @@ -151,8 +159,6 @@ def tearDown(self): self.patcher_rank.stop() def test_build_torch_profiler_returns_active_handle(self): - import tempfile - with tempfile.TemporaryDirectory() as tmpdir: profiler = Profiler( Profiler.Config( @@ -167,6 +173,126 @@ def test_build_torch_profiler_returns_active_handle(self): with profiler: self.assertIsNotNone(profiler.torch_profiler) + @mock.patch("torchtitan.tools.profiler.device_module.memory._record_memory_history") + def test_context_reuses_started_memory_profiler(self, mock_record_history): + with tempfile.TemporaryDirectory() as tmpdir: + profiler = Profiler( + Profiler.Config(enable_memory_snapshot=True), + global_step=0, + base_folder=tmpdir, + ) + memory_profiler = profiler.start_memory_profiler() + with profiler as active_profiler: + self.assertIs(active_profiler.memory_profiler, memory_profiler) + + mock_record_history.assert_called_once() + + +class _FakeProfileMode(enum.Enum): + RUNTIME = 1 + CPU_UTIL = 2 + HOST_MEMORY = 3 + DEVICE = 4 + + +class _FakeNeuronConfig: + __slots__ = ("kwargs",) + + def __init__(self, **kwargs): + self.kwargs = kwargs + + +class _FakeNeuronProfiler: + def __init__(self, config): + self.config = config + + def export_trace(self, prof): # matches torch_neuronx NeuronProfiler API + pass + + +def _fake_torch_neuronx_profiling(): + mod = types.ModuleType("torch_neuronx.profiling") + mod.NeuronConfig = _FakeNeuronConfig + mod.NeuronProfiler = _FakeNeuronProfiler + mod.ProfileMode = _FakeProfileMode + return mod + + +class TestNeuronProfiling(unittest.TestCase): + """Tests for the Neuron branch of build_torch_profiler. + + _enable_neuron_profiling is host-independent (it never touches real Neuron + hardware -- it only wires torch_neuronx.profiling into the torch profiler), so + these run on any CPU host with torch_neuronx.profiling mocked. + """ + + def _call(self, env): + profiler = Profiler(Profiler.Config(enable_profiling=True)) + activities = [torch.profiler.ProfilerActivity.CPU] + profile_kwargs: dict = {} + with mock.patch.dict( + sys.modules, {"torch_neuronx.profiling": _fake_torch_neuronx_profiling()} + ), mock.patch.dict(os.environ, env, clear=False): + os.environ.pop("NEURON_PROFILE_DEVICE", None) + os.environ.pop("NEURON_PROFILE_NEFF_CACHE_DIR", None) + os.environ.pop("NEURON_PROFILE_CAPTURE_NC", None) + os.environ.update(env) + handler = profiler._enable_neuron_profiling( + activities, profile_kwargs, "/tmp/traces" + ) + return handler, activities, profile_kwargs + + def test_system_modes_default(self): + """Without NEURON_PROFILE_DEVICE, only system + host modes are enabled.""" + handler, activities, kwargs = self._call(env={}) + self.assertIsNotNone(handler) + self.assertIn(torch.profiler.ProfilerActivity.PrivateUse1, activities) + self.assertIn("experimental_config", kwargs) + self.assertTrue(kwargs["with_stack"]) + cfg = kwargs["experimental_config"] + self.assertEqual( + {m.name for m in cfg.kwargs["modes"]}, + {"RUNTIME", "CPU_UTIL", "HOST_MEMORY"}, + ) + self.assertEqual(cfg.kwargs["profile_output_dir"], "/tmp/traces") + self.assertEqual(cfg.kwargs["max_events_per_nc"], 10_000_000) + # Defaults to torch_neuronx's compile cache so DEVICE .ntff resolves NEFFs. + self.assertEqual(cfg.kwargs["neff_cache_dir"], "/tmp/neff_cache") + self.assertEqual(cfg.kwargs["capture_enabled_for_nc"], "0,1") + + def test_device_mode_enabled_by_env(self): + """NEURON_PROFILE_DEVICE=1 adds the heavy DEVICE mode.""" + _, _, kwargs = self._call(env={"NEURON_PROFILE_DEVICE": "1"}) + cfg = kwargs["experimental_config"] + self.assertIn("DEVICE", {m.name for m in cfg.kwargs["modes"]}) + + def test_neff_cache_dir_env_override(self): + """NEURON_PROFILE_NEFF_CACHE_DIR overrides the default neff cache dir.""" + _, _, kwargs = self._call( + env={"NEURON_PROFILE_NEFF_CACHE_DIR": "/custom/neffs"} + ) + cfg = kwargs["experimental_config"] + self.assertEqual(cfg.kwargs["neff_cache_dir"], "/custom/neffs") + + def test_capture_nc_env_override(self): + """NEURON_PROFILE_CAPTURE_NC overrides which NeuronCores are captured.""" + _, _, kwargs = self._call(env={"NEURON_PROFILE_CAPTURE_NC": "0-63"}) + cfg = kwargs["experimental_config"] + self.assertEqual(cfg.kwargs["capture_enabled_for_nc"], "0-63") + + def test_import_failure_falls_back_to_cpu_only(self): + """When torch_neuronx.profiling is unavailable, returns None untouched.""" + profiler = Profiler(Profiler.Config(enable_profiling=True)) + activities = [torch.profiler.ProfilerActivity.CPU] + profile_kwargs: dict = {} + with mock.patch.dict(sys.modules, {"torch_neuronx.profiling": None}): + handler = profiler._enable_neuron_profiling( + activities, profile_kwargs, "/tmp/traces" + ) + self.assertIsNone(handler) + self.assertEqual(activities, [torch.profiler.ProfilerActivity.CPU]) + self.assertEqual(profile_kwargs, {}) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/test_quantization.py b/tests/unit_tests/test_quantization.py index 5a21e6a661e..dfea44ce8b2 100644 --- a/tests/unit_tests/test_quantization.py +++ b/tests/unit_tests/test_quantization.py @@ -1,8 +1,10 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. + import pytest from torchtitan.components.quantization import Float8Linear diff --git a/tests/unit_tests/test_rope.py b/tests/unit_tests/test_rope.py index f97c02bec01..30974dac6ff 100644 --- a/tests/unit_tests/test_rope.py +++ b/tests/unit_tests/test_rope.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import dataclasses import unittest @@ -20,7 +21,14 @@ CosSinRoPE, RoPE, ) -from torchtitan.models.qwen3_5.rope import MRoPE + +# qwen3_5 pulls in `fla` (flash-linear-attention), an optional dep not present +# in every environment (e.g. the Neuron stack). Guard the import so only the +# MRoPE test is skipped when it is missing, not the whole module. +try: + from torchtitan.models.qwen3_5.rope import MRoPE +except ModuleNotFoundError: + MRoPE = None class TestApplyRotaryEmbCosSin(unittest.TestCase): @@ -163,6 +171,9 @@ def test_out_of_range_positions_raises(self): self.rope(xq, xk, positions) +@unittest.skipIf( + MRoPE is None, "requires 'fla' (flash-linear-attention) via the qwen3_5 model" +) class TestMRoPECache(unittest.TestCase): def test_forward_accepts_three_axis_positions(self): torch.manual_seed(42) @@ -236,6 +247,54 @@ def test_decoder_builds_distinct_rope_configs_per_attention_layer(self): ) +class TestPrefixCache(unittest.TestCase): + def test_prefix_matches_sequential_gather_outputs_and_gradients(self): + torch.manual_seed(42) + for rope_cls in (ComplexRoPE, CosSinRoPE): + with self.subTest(rope=rope_cls.__name__): + config = rope_cls.Config(dim=8, max_seq_len=64) + gather = config.build() + prefix = dataclasses.replace( + config, use_prefix_cache_for_positions=True + ).build() + query = torch.randn(2, 8, 4, 8, requires_grad=True) + key = torch.randn(2, 8, 2, 8, requires_grad=True) + query_prefix = query.detach().clone().requires_grad_() + key_prefix = key.detach().clone().requires_grad_() + positions = torch.arange(8).expand(2, -1) + expected = gather(query, key, positions) + actual = prefix(query_prefix, key_prefix, positions) + for output, reference in zip(actual, expected): + torch.testing.assert_close(output, reference, rtol=0, atol=0) + sum(output.sum() for output in expected).backward() + sum(output.sum() for output in actual).backward() + torch.testing.assert_close( + query_prefix.grad, query.grad, rtol=0, atol=0 + ) + torch.testing.assert_close(key_prefix.grad, key.grad, rtol=0, atol=0) + + def test_gather_preserves_reset_and_offset_positions(self): + query = torch.ones(2, 4, 2, 8) + key = torch.ones(2, 4, 1, 8) + positions = torch.tensor([[10, 11, 12, 13], [0, 1, 0, 1]]) + for rope_cls in (ComplexRoPE, CosSinRoPE): + with self.subTest(rope=rope_cls.__name__): + config = rope_cls.Config(dim=8, max_seq_len=64) + rope = config.build() + selected_cache = rope.cache[positions].unsqueeze(2) + expected = rope_cls.apply_rotary_emb(query, key, selected_cache) + actual = rope(query, key, positions) + prefix = dataclasses.replace( + config, use_prefix_cache_for_positions=True + ).build() + prefix_outputs = prefix(query, key, positions) + for output, reference, prefix_output in zip( + actual, expected, prefix_outputs + ): + torch.testing.assert_close(output, reference, rtol=0, atol=0) + self.assertFalse(torch.equal(output, prefix_output)) + + class TestUpdateFromConfigSeqLenValidation(unittest.TestCase): """update_from_config must reject seq_len > rope.max_seq_len.""" diff --git a/tests/unit_tests/test_trainer_profiler_ordering.py b/tests/unit_tests/test_trainer_profiler_ordering.py new file mode 100644 index 00000000000..66ce13ec72a --- /dev/null +++ b/tests/unit_tests/test_trainer_profiler_ordering.py @@ -0,0 +1,147 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +import contextlib +import sys +import types +import unittest +from types import SimpleNamespace +from unittest import mock + +sys.modules.setdefault( + "spmd_types", + types.SimpleNamespace(no_typecheck=contextlib.nullcontext), +) + +try: + from torchtitan.trainer import Trainer +except ImportError as error: + Trainer = None + _TRAINER_IMPORT_ERROR = error +else: + _TRAINER_IMPORT_ERROR = None + + +class _FakeProfiler: + def __init__(self, events): + self.events = events + self.memory_profiler_started = False + self.start_memory_profiler_call_count = 0 + + def start_memory_profiler(self): + self.start_memory_profiler_call_count += 1 + if not self.memory_profiler_started: + self.events.append("start_memory_profiler") + self.memory_profiler_started = True + + def set_memory_profiler_step(self, global_step): + self.events.append(("set_memory_profiler_step", global_step)) + + def dump_memory_snapshot(self, name): + self.events.append(("dump_memory_snapshot", name)) + + def active(self, *, global_step=0, base_folder="", leaf_folder=""): + self.events.append(("active", global_step, base_folder, leaf_folder)) + return self + + def __enter__(self): + self.events.append("enter_profiler_context") + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.events.append("exit_profiler_context") + return False + + def step(self): + self.events.append("profiler_step") + + +class _FakeProfilerConfig: + def __init__(self, events): + self.events = events + self.enable_memory_snapshot = True + self.profiler = _FakeProfiler(events) + + def build(self, *, global_step=0, base_folder="", leaf_folder=""): + self.events.append(("build_profiler", global_step, base_folder, leaf_folder)) + return self.profiler + + +class _FakeCheckpointer: + def __init__(self, events, trainer, loaded_step, loaded=True): + self.events = events + self.trainer = trainer + self.loaded_step = loaded_step + self.loaded = loaded + + def load(self, *, step): + self.events.append(("checkpoint_load", step)) + self.trainer.step = self.loaded_step + return self.loaded + + +@unittest.skipIf( + Trainer is None, f"Trainer import unavailable: {_TRAINER_IMPORT_ERROR}" +) +class TestTrainerProfilerOrdering(unittest.TestCase): + def _trainer(self, *, checkpoint_loaded=True): + events = [] + trainer = object.__new__(Trainer) + trainer.step = 0 + trainer.config = SimpleNamespace( + checkpoint=SimpleNamespace(load_step=-1), + dump_folder="/tmp/torchtitan-test-output", + profiler=_FakeProfilerConfig(events), + ) + trainer.checkpointer = _FakeCheckpointer( + events, trainer, loaded_step=123, loaded=checkpoint_loaded + ) + trainer.profiler = trainer.config.profiler.build( + global_step=0, + base_folder="/tmp/torchtitan-test-output", + ) + trainer.profiler.start_memory_profiler() + trainer.dataloader = object() + trainer.batch_generator = mock.Mock(return_value=iter(())) + trainer.should_continue_training = mock.Mock(return_value=False) + return trainer, events + + @mock.patch("torch.distributed.get_rank", return_value=1) + @mock.patch("torchtitan.trainer.sl.log_trace_instant") + def test_memory_history_starts_before_checkpoint_load(self, *_): + trainer, events = self._trainer(checkpoint_loaded=True) + + Trainer.train(trainer) + + self.assertEqual(trainer.profiler.start_memory_profiler_call_count, 1) + self.assertEqual( + events, + [ + ("build_profiler", 0, "/tmp/torchtitan-test-output", ""), + "start_memory_profiler", + ("active", 0, "/tmp/torchtitan-test-output", ""), + ("checkpoint_load", -1), + ("set_memory_profiler_step", 123), + ("active", 123, "/tmp/torchtitan-test-output", ""), + "enter_profiler_context", + "exit_profiler_context", + ], + ) + + @mock.patch("torch.distributed.get_rank", return_value=1) + @mock.patch("torchtitan.trainer.sl.log_trace_instant") + def test_train_uses_existing_profiler_when_checkpoint_not_loaded(self, *_): + trainer, events = self._trainer(checkpoint_loaded=False) + + Trainer.train(trainer) + + self.assertEqual( + [event for event in events if event[0] == "build_profiler"], + [("build_profiler", 0, "/tmp/torchtitan-test-output", "")], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/__init__.py b/torchtitan/__init__.py index a881d3e72a2..1f4a4275f26 100644 --- a/torchtitan/__init__.py +++ b/torchtitan/__init__.py @@ -7,6 +7,6 @@ from importlib.metadata import version try: - __version__ = version("torchtitan") + __version__ = version("torchtitan-neuron") except Exception: __version__ = "0.0.0+unknown" diff --git a/torchtitan/components/checkpoint.py b/torchtitan/components/checkpoint.py index 6581955eee6..0120f14d735 100644 --- a/torchtitan/components/checkpoint.py +++ b/torchtitan/components/checkpoint.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from __future__ import annotations @@ -21,10 +22,20 @@ import torch.distributed as dist import torch.distributed.checkpoint as dcp import torch.nn as nn -from torch.distributed.checkpoint import HuggingFaceStorageWriter +from torch.distributed.checkpoint import DefaultLoadPlanner, HuggingFaceStorageWriter from torch.distributed.checkpoint._consolidate_hf_safetensors import ( consolidate_safetensors_files_on_every_rank, ) +from torch.distributed.checkpoint.default_planner import create_default_local_load_plan +from torch.distributed.checkpoint.metadata import ( + ChunkStorageMetadata, + Metadata, + TensorStorageMetadata, +) +from torch.distributed.checkpoint.planner import LoadPlan, ReadItem +from torch.distributed.checkpoint.planner_helpers import ( + create_read_items_for_chunk_list, +) from torch.distributed.checkpoint.staging import DefaultStager, StagingOptions from torch.distributed.checkpoint.state_dict_saver import ( AsyncCheckpointerType, @@ -37,6 +48,7 @@ from torchtitan.components.optimizer import OptimizersContainer from torchtitan.config import Configurable, TORCH_DTYPE_MAP from torchtitan.observability import structured_logger as sl +from torchtitan.protocols.checkpoint_load import FusedQKVLoadTarget, plan_fused_qkv_load from torchtitan.protocols.state_dict_adapter import BaseStateDictAdapter from torchtitan.tools.logging import logger from torchtitan.tools.utils import GarbageCollection @@ -97,15 +109,22 @@ class ModelWrapper(Stateful): cache will become stale. """ - def __init__(self, model: nn.Module | list[nn.Module]) -> None: + def __init__( + self, model: nn.Module | list[nn.Module], *, cache_state_dict: bool = True + ) -> None: self.model = [model] if isinstance(model, nn.Module) else model - self.cached_state_dict = self._get_state_dict() + self.cache_state_dict = cache_state_dict + self.cached_state_dict = self._get_state_dict() if cache_state_dict else None def _get_state_dict(self) -> dict[str, Any]: # TorchTitan already makes model state_dict keys canonical. return {k: v for model in self.model for k, v in model.state_dict().items()} def state_dict(self) -> dict[str, Any]: + if not self.cache_state_dict: + return self._get_state_dict() + + assert self.cached_state_dict is not None # Recompute the state dict so hook-produced tensors reflect the current # parameters, then merge into the cache without changing storage objects. for key, value in self._get_state_dict().items(): @@ -126,7 +145,211 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None: for model in self.model: model.load_state_dict(state_dict, strict=False) # Refresh the cache so state_dict() reflects the freshly loaded values. - self.cached_state_dict = self._get_state_dict() + self.cached_state_dict = ( + self._get_state_dict() if self.cache_state_dict else None + ) + + +class FusedQKVLoadPlanner(DefaultLoadPlanner): + """Stream checkpoint Q/K/V tensors directly into their fused runtime tensor.""" + + def __init__(self) -> None: + super().__init__(flatten_state_dict=False, flatten_sharded_tensors=False) + self.targets: dict[str, FusedQKVLoadTarget] = {} + + def set_up_planner( + self, + state_dict: dict[str, Any], + metadata: Metadata | None = None, + is_coordinator: bool = False, + ) -> None: + # Do not call DefaultLoadPlanner.set_up_planner(): it materializes meta + # tensors before reading. FusedQKVLoadTarget deliberately carries only + # shape/layout metadata so all Q/K/V device buffers are not allocated at + # once before DCP starts streaming reads. + self.original_state_dict = state_dict + self.state_dict = state_dict + self.metadata = metadata + self.is_coordinator = is_coordinator + self.targets = { + key: value + for key, value in state_dict.items() + if isinstance(value, FusedQKVLoadTarget) + } + + def create_local_plan(self) -> LoadPlan: + if self.metadata is None: + raise AssertionError("FusedQKVLoadPlanner metadata is not initialized.") + + regular_state_dict = { + key: value + for key, value in self.state_dict.items() + if not isinstance(value, FusedQKVLoadTarget) + } + regular_plan = create_default_local_load_plan( + regular_state_dict, self.metadata, strict=True + ) + requests = list(regular_plan.items) + + for fqn, target in self.targets.items(): + if fqn not in self.metadata.state_dict_metadata: + raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.") + checkpoint_metadata = self.metadata.state_dict_metadata[fqn] + if not isinstance(checkpoint_metadata, TensorStorageMetadata): + raise ValueError( + f"Expected tensor checkpoint metadata for {fqn}, " + f"found {type(checkpoint_metadata)}." + ) + if checkpoint_metadata.size != target.shape: + raise ValueError( + f"Size mismatch between saved {checkpoint_metadata.size} and " + f"current: {target.shape} for {fqn}" + ) + + requests.extend( + create_read_items_for_chunk_list( + fqn, + checkpoint_metadata, + [ + ChunkStorageMetadata( + offsets=torch.Size([0] * len(target.shape)), + sizes=target.shape, + ) + ], + ) + ) + + return LoadPlan(requests) + + def resolve_tensor(self, read_item: ReadItem) -> torch.Tensor: + target = self.targets.get(read_item.dest_index.fqn) + if target is None: + return super().resolve_tensor(read_item) + + fused_tensor = target.fused_tensor + local_tensor = ( + fused_tensor._local_tensor + if isinstance(fused_tensor, DTensor) + else fused_tensor + ) + return torch.empty( + read_item.lengths, + dtype=target.dtype, + device=local_tensor.device, + ) + + @staticmethod + def _local_fused_tensor( + target: FusedQKVLoadTarget, + ) -> tuple[torch.Tensor, torch.Size, torch.Size]: + fused_tensor = target.fused_tensor + if isinstance(fused_tensor, DTensor): + chunks = fused_tensor.__create_chunk_list__() + if len(chunks) != 1: + raise RuntimeError( + "FusedQKVLoadPlanner expects one local DTensor chunk, " + f"found {len(chunks)}." + ) + chunk = chunks[0] + return fused_tensor._local_tensor, chunk.offsets, chunk.sizes + + shape = fused_tensor.shape + return ( + fused_tensor, + torch.Size([0] * len(shape)), + torch.Size(shape), + ) + + @staticmethod + def _copy_projection_chunk( + target: FusedQKVLoadTarget, + read_item: ReadItem, + tensor: torch.Tensor, + ) -> None: + ( + local_fused, + local_offsets, + local_sizes, + ) = FusedQKVLoadPlanner._local_fused_tensor(target) + ndim = len(target.shape) + chunk_offsets = read_item.dest_index.offset or torch.Size([0] * ndim) + source_offsets = torch.Size( + chunk_offsets[dim] + read_item.dest_offsets[dim] for dim in range(ndim) + ) + + source_row_start = source_offsets[0] + source_row_end = source_row_start + read_item.lengths[0] + local_row_start = local_offsets[0] + local_row_end = local_row_start + local_sizes[0] + + if ndim == 2: + source_col_start = source_offsets[1] + source_col_end = source_col_start + read_item.lengths[1] + local_col_start = local_offsets[1] + local_col_end = local_col_start + local_sizes[1] + copy_col_start = max(source_col_start, local_col_start) + copy_col_end = min(source_col_end, local_col_end) + if copy_col_start >= copy_col_end: + return + else: + source_col_start = local_col_start = copy_col_start = 0 + copy_col_end = 0 + + projection_rows = ( + target.heads_per_kv * target.head_dim + if target.projection == "q" + else target.head_dim + ) + projection_fused_offset = { + "q": 0, + "k": target.heads_per_kv * target.head_dim, + "v": (target.heads_per_kv + 1) * target.head_dim, + }[target.projection] + fused_group_rows = (target.heads_per_kv + 2) * target.head_dim + + for group in range(target.n_kv_heads): + projection_group_start = group * projection_rows + projection_group_end = projection_group_start + projection_rows + copy_source_start = max(source_row_start, projection_group_start) + copy_source_end = min(source_row_end, projection_group_end) + if copy_source_start >= copy_source_end: + continue + + fused_start = ( + group * fused_group_rows + + projection_fused_offset + + copy_source_start + - projection_group_start + ) + fused_end = fused_start + copy_source_end - copy_source_start + + clipped_fused_start = max(fused_start, local_row_start) + clipped_fused_end = min(fused_end, local_row_end) + if clipped_fused_start >= clipped_fused_end: + continue + + clipped_source_start = copy_source_start + clipped_fused_start - fused_start + row_count = clipped_fused_end - clipped_fused_start + source_row_offset = clipped_source_start - source_row_start + local_row_offset = clipped_fused_start - local_row_start + + source_view = tensor.narrow(0, source_row_offset, row_count) + target_view = local_fused.narrow(0, local_row_offset, row_count) + if ndim == 2: + col_count = copy_col_end - copy_col_start + source_view = source_view.narrow( + 1, copy_col_start - source_col_start, col_count + ) + target_view = target_view.narrow( + 1, copy_col_start - local_col_start, col_count + ) + target_view.copy_(source_view) + + def commit_tensor(self, read_item: ReadItem, tensor: torch.Tensor) -> None: + target = self.targets.get(read_item.dest_index.fqn) + if target is None: + return super().commit_tensor(read_item, tensor) + self._copy_projection_chunk(target, read_item, tensor) class Terminate: @@ -372,6 +595,15 @@ class Config(Configurable.Config): checkpoints to be loaded without saving any during the training. """ + load_model_state_dict_with_streaming_planner: bool = False + """ + Load native DCP model-only checkpoints with one DCP call while streaming + checkpoint-facing Q/K/V tensors directly into FusedQKVLinear's fused + parameter. Requires load_only=True so ModelWrapper does not materialize + checkpoint-facing Q/K/V tensors for its save-oriented state-dict cache. + Other checkpoint formats continue to use the standard load path. + """ + def __post_init__(self): if not self.folder.strip(): raise ValueError("The 'folder' field cannot be empty.") @@ -404,6 +636,11 @@ def __post_init__(self): ) if self.last_save_in_hf and not self.last_save_model_only: raise ValueError("last_save_in_hf requires last_save_model_only=True.") + if self.load_model_state_dict_with_streaming_planner and not self.load_only: + raise ValueError( + "load_model_state_dict_with_streaming_planner requires " + "load_only=True." + ) async_lowered = self.async_mode.lower() if async_lowered in ("disabled", "async", "async_with_pinned_mem"): @@ -445,7 +682,7 @@ def __init__( self.states = states self.states.update( { - MODEL: ModelWrapper(model_parts), + MODEL: ModelWrapper(model_parts, cache_state_dict=not config.load_only), OPTIMIZER: optimizers, DATALOADER: dataloader, LR_SCHEDULER: lr_schedulers, @@ -454,6 +691,9 @@ def __init__( # Loading & Saving Policy self.load_only = config.load_only + self.load_model_state_dict_with_streaming_planner = ( + config.load_model_state_dict_with_streaming_planner + ) self.exclude_from_loading = config.exclude_from_loading self.initial_load_path = config.initial_load_path self.initial_load_model_only = config.initial_load_model_only @@ -670,6 +910,48 @@ def dcp_load( if MODEL in self.states: self.states[MODEL].load_state_dict(state_dict) + def dcp_load_model_state_dict_with_streaming_planner( + self, checkpoint_id: str + ) -> bool: + """Load fused QKV checkpoint tensors with one DCP planner invocation.""" + wrapper = cast(ModelWrapper, self.states[MODEL]) + with plan_fused_qkv_load(): + state_dict = { + key: value + for model in wrapper.model + for key, value in model.state_dict().items() + } + + if not any( + isinstance(value, FusedQKVLoadTarget) for value in state_dict.values() + ): + logger.warning( + "Streaming checkpoint load found no FusedQKVLinear targets; " + "falling back to standard checkpoint loading." + ) + return False + + logger.info("Loading model checkpoint with the fused-QKV streaming planner.") + dcp.load( + state_dict, + checkpoint_id=checkpoint_id, + planner=FusedQKVLoadPlanner(), + ) + + # Ordinary entries either alias live parameters or need their normal module + # load hook. Remove Q/K/V metadata targets so FusedQKVLinear's merge hook + # cannot allocate the cat buffer that the planner intentionally avoids. + state_dict = { + key: value + for key, value in state_dict.items() + if not isinstance(value, FusedQKVLoadTarget) + } + for model in wrapper.model: + model.load_state_dict(state_dict, strict=False) + + logger.info("Finished fused-QKV streaming model checkpoint load.") + return True + @sl.log_trace_span("checkpoint_save") @torch.no_grad() def save(self, curr_step: int, last_step: bool = False) -> bool: @@ -864,13 +1146,26 @@ def load(self, step: int = -1) -> bool: logger.info(f"Loading the checkpoint from {checkpoint_id}.") begin = time.monotonic() - states = self._states_to_load(model_only) - self.dcp_load( - states, - checkpoint_id=checkpoint_id, - from_hf=from_hf, - from_quantized=from_quantized, - ) + loaded_with_streaming_planner = False + if self.load_model_state_dict_with_streaming_planner and model_only: + if from_hf or from_quantized: + logger.warning( + "checkpoint.load_model_state_dict_with_streaming_planner is only " + "supported for native DCP loads; falling back to standard load." + ) + else: + loaded_with_streaming_planner = ( + self.dcp_load_model_state_dict_with_streaming_planner(checkpoint_id) + ) + + if not loaded_with_streaming_planner: + states = self._states_to_load(model_only) + self.dcp_load( + states, + checkpoint_id=checkpoint_id, + from_hf=from_hf, + from_quantized=from_quantized, + ) GarbageCollection.collect("GC collection for checkpoint loading.") logger.info( diff --git a/torchtitan/components/metrics.py b/torchtitan/components/metrics.py index f63612eabac..c4370662498 100644 --- a/torchtitan/components/metrics.py +++ b/torchtitan/components/metrics.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import os import time @@ -257,6 +258,27 @@ def _get_metrics_rank( return (world_size // pp_size) * (pp_size - 1) +def compute_mfu( + num_flops_per_token: int, + tps: float, + gpu_peak_flops: float, + has_quantization: bool, +) -> float | None: + """Model FLOPS utilization (%), or None when it would be misleading. + + MFU is defined against the BF16 peak FLOPS (PaLM paper, + https://arxiv.org/abs/2204.02311). When quantization (FP8/MX) is active the + BF16 peak understates the effective peak, so MFU is skipped -- except on + Neuron trn3 MXFP8, where the tensor-engine peak is a known 4x the BF16 peak, + so we report MFU against ``4 * gpu_peak_flops``. + """ + if has_quantization: + if utils.has_neuron_capability("mxfp8"): + return 100 * num_flops_per_token * tps / (4 * gpu_peak_flops) + return None + return 100 * num_flops_per_token * tps / gpu_peak_flops + + class MetricsProcessor(Configurable): """Metrics processor to processes the metrics and log metrics. @@ -477,16 +499,11 @@ def log( tps = self.ntokens_since_last_log / ( time_delta * self.parallel_dims.non_data_parallel_size ) - # model FLOPS utilization - # For its definition and calculation, please refer to the PaLM paper: - # https://arxiv.org/abs/2204.02311 - # MFU is based on BF16 peak FLOPS which is misleading when quantization - # (FP8/MX) is active, so we skip it in that case. + # model FLOPS utilization (see compute_mfu for the quantization handling) tflops = self.num_flops_per_token * tps / 1e12 - if self.has_quantization: - mfu = None - else: - mfu = 100 * self.num_flops_per_token * tps / self.gpu_peak_flops + mfu = compute_mfu( + self.num_flops_per_token, tps, self.gpu_peak_flops, self.has_quantization + ) time_end_to_end = time_delta / self.config.log_freq time_data_loading = sum(self.data_loading_times) / len(self.data_loading_times) diff --git a/torchtitan/components/optimizer.py b/torchtitan/components/optimizer.py index 1a7afaf9a69..8c0af1e829b 100644 --- a/torchtitan/components/optimizer.py +++ b/torchtitan/components/optimizer.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import re from collections import defaultdict @@ -409,7 +410,7 @@ def _iter_moe_layers( if getattr(transformer_block, "moe_enabled", False): yield transformer_block, cast(_MoELike, transformer_block.moe) - def _should_register_moe_balancing_hook(model_parts: list[nn.Module]) -> bool: + def _should_register_expert_bias_hook(model_parts: list[nn.Module]) -> bool: moe_layers = list(_iter_moe_layers(model_parts)) if not moe_layers: return False @@ -422,7 +423,15 @@ def _should_register_moe_balancing_hook(model_parts: list[nn.Module]) -> bool: "across all MoE layers. Either set it for every MoE layer " "or leave it unset for all MoE layers." ) - return load_balance_enabled + mechanisms = { + getattr(moe, "load_balancing", "expert_bias") for _, moe in moe_layers + } + if len(mechanisms) > 1: + raise ValueError( + "MoE load_balancing must be the same across all MoE layers" + ) + # The global-batch loss has no expert bias to update. + return load_balance_enabled and mechanisms == {"expert_bias"} # for MoE auxiliary-loss-free load balancing def _is_recomputation_enabled(module): @@ -520,7 +529,7 @@ def _update_expert_bias( moe.expert_bias_E.add_(expert_bias_delta_E) moe.tokens_per_expert_E.zero_() - if _should_register_moe_balancing_hook(model_parts): + if _should_register_expert_bias_hook(model_parts): optimizers.register_step_pre_hook( lambda *args, **kwargs: _update_expert_bias( model_parts, parallel_dims=parallel_dims diff --git a/torchtitan/components/quantization/mx.py b/torchtitan/components/quantization/mx.py index c039e289b48..5bc5309a593 100644 --- a/torchtitan/components/quantization/mx.py +++ b/torchtitan/components/quantization/mx.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from dataclasses import dataclass, field, fields from importlib.util import find_spec @@ -13,7 +14,7 @@ from torchtitan.models.common.nn_modules import Linear from torchtitan.protocols.module import Module from torchtitan.tools.logging import logger -from torchtitan.tools.utils import has_cuda_capability +from torchtitan.tools.utils import has_cuda_capability, has_neuron_capability from .utils import swap_token_dispatcher @@ -33,14 +34,33 @@ class MXFP8Linear(TorchAOMXFP8Linear, Module): class Config(Linear.Config): """Drop-in replacement for Linear.Config that builds MXFP8Linear.""" - pass + kernel_preference: str = "auto" + """ + torchao ``KernelPreference`` name for the MXFP8 GEMMs. Kept as a + plain string so this config tree imports without torchao installed. + Resolution of ``"auto"`` is device-gated in ``__init__``. + """ def __init__(self, config: Config): + from torchao.quantization.quantize_.common.kernel_preference import ( + KernelPreference, + ) + + # "auto" resolves to the Neuron kernel on trn3, else stays AUTO. An + # explicit preference is honored as-is. + pref = config.kernel_preference + if pref == "auto" and has_neuron_capability("mxfp8"): + pref = "neuron" + + if pref == "neuron" and not hasattr(KernelPreference, "NEURON"): + raise ValueError("missing Neuron torchao dependency") + TorchAOMXFP8Linear.__init__( self, config.in_features, config.out_features, bias=config.bias, + kernel_preference=KernelPreference(pref), ) except ImportError: @@ -58,6 +78,11 @@ class Config(QuantizationConverter.Config): Only Linear.Config entries whose FQN contains a match are converted. If empty, all Linear modules are converted. """ + kernel_preference: str = "auto" + """ + torchao ``KernelPreference`` name forwarded to each converted + ``MXFP8Linear.Config`` (resolved in ``MXFP8Linear.__init__``). + """ def __init__(self, config: Config): self.config = config @@ -67,8 +92,10 @@ def __init__(self, config: Config): "torchao is not installed. Please install it to use MXFP8 linear layers." ) - if not has_cuda_capability(10, 0): - raise ValueError("MXFP8 is only supported on SM100 or later architectures") + if not (has_cuda_capability(10, 0) or has_neuron_capability("mxfp8")): + raise ValueError( + "MXFP8 is only supported on CUDA SM100 or later, or Neuron trn3 or later" + ) if not self.config.model_compile_enabled: logger.warning( @@ -86,6 +113,7 @@ def convert(self, model_config): out_features=config.out_features, bias=config.bias, param_init=config.param_init, + kernel_preference=self.config.kernel_preference, ) if parent is None: model_config = new_config diff --git a/torchtitan/distributed/activation_checkpoint.py b/torchtitan/distributed/activation_checkpoint.py index 476a209adf3..6dacca65d8c 100644 --- a/torchtitan/distributed/activation_checkpoint.py +++ b/torchtitan/distributed/activation_checkpoint.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. # This file provides the util functions to apply activation checkpointing to the model. # Technically, this is not a part of distributed, but distributed module is the best place to put it. @@ -48,6 +49,14 @@ def _get_default_save_ops() -> set: # FlexAttention (torch.ops.higher_order.flex_attention is the same object) torch._higher_order_ops.flex_attention, torch.ops.aten.linear.default, + # Neuron MXFP8: these replace aten.linear, so they need the same save treatment. + # Without them the backward re-runs the forward, which re-does the encoding + # all-gathers. + # Tuple form: silently skipped off Neuron. + (torch.ops, "neuron._scaled_mm.default"), + (torch.ops, "neuron.fused_quantize_mx_scaled_mm.default"), + (torch.ops, "neuron.prequantized_mx_scaled_mm_nki.default"), + (torch.ops, "neuron.fused_quantize_mx_scaled_mm_nki.default"), # topk can be non-deterministic; save to keep MoE expert assignments # stable between forward and recompute. torch.ops.aten.topk.default, @@ -105,7 +114,6 @@ def _disable_dynamo_lru_cache() -> None: # here is to disable the LRU cache, and select graphs in insertion order instead. # # Also see: https://github.com/pytorch/pytorch/issues/166926 - # pyrefly: ignore [missing-attribute] torch._C._dynamo.eval_frame._set_lru_cache(False) @@ -187,8 +195,9 @@ class SelectiveAC(ActivationCheckpointing): Saves the outputs of compute/communication ops that are expensive to recompute (see ``get_save_ops``) while recomputing the rest, and recomputes - every second matmul to balance memory and compute. Override ``get_save_ops`` - in a subclass to tune which ops are saved. + every ``matmul_recompute_freq``-th matmul (default: every second) to balance + memory and compute. Override ``get_save_ops`` in a subclass to tune which + ops are saved. """ @dataclass(kw_only=True, slots=True) @@ -206,6 +215,17 @@ class Config(ActivationCheckpointing.Config): ANY mm with shape matching (*, in) x (in, out) will be force recomputed. """ + matmul_recompute_freq: int = 2 + """ + Recompute one matmul/linear out of every ``matmul_recompute_freq`` under + the default SAC policy (the rest are saved). ``2`` (default) recomputes + every second matmul, matching the historical policy; larger values + recompute fewer matmuls (save more); ``0`` saves every matmul (recompute + none) -- used by memory-headroom configs that need all mms saved. + Matmuls matched by ``force_recompute_mm_shapes_by_fqns`` are always + recomputed regardless of this value. Must be >= 0. + """ + def get_save_ops(self) -> set: """Returns the set of ops whose activations should be saved. Override to customize the save set.""" @@ -215,6 +235,11 @@ def _wrap_block( self, module: nn.Module, *, base_fqn: str | None = None ) -> nn.Module: config = cast("SelectiveAC.Config", self.config) + if config.matmul_recompute_freq < 0: + raise ValueError( + "matmul_recompute_freq must be >= 0, got " + f"{config.matmul_recompute_freq}" + ) save_ops = self.get_save_ops() # Collect weight shapes to force-recompute, stored as mm RHS shape @@ -239,6 +264,10 @@ def _wrap_block( # Some backends (e.g. PrivateUse1) register aten.linear as a leaf op # instead of decomposing it into aten.mm, so we must handle both. mm_ops = (torch.ops.aten.mm.default, torch.ops.aten.linear.default) + # The policy closure captures only plain values: torch.compile precompile + # pickles what the guard state reaches, and a nested config class cannot be + # pickled by reference. + mm_recompute_freq = config.matmul_recompute_freq def _get_custom_policy(): meta = {"forward_mm_count": 0, "recompute_mm_count": 0} @@ -266,9 +295,14 @@ def wrapped_policy(ctx, func, *args, **kwargs) -> CheckpointPolicy: return CheckpointPolicy.PREFER_RECOMPUTE meta[mm_count_key] += 1 - # Save all compute/comm ops, except every second mm/linear. + # Save all compute/comm ops, except every + # matmul_recompute_freq-th mm/linear (freq=0 saves every mm). if func in save_ops: - if func in mm_ops and meta[mm_count_key] % 2 == 0: + if ( + func in mm_ops + and mm_recompute_freq != 0 + and meta[mm_count_key] % mm_recompute_freq == 0 + ): return CheckpointPolicy.PREFER_RECOMPUTE return CheckpointPolicy.MUST_SAVE return CheckpointPolicy.PREFER_RECOMPUTE diff --git a/torchtitan/distributed/compile.py b/torchtitan/distributed/compile.py index b57cc036464..1303b6a6898 100644 --- a/torchtitan/distributed/compile.py +++ b/torchtitan/distributed/compile.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import contextlib from collections.abc import Callable @@ -45,9 +46,7 @@ def apply_compile(model: nn.Module, compile_config: CompileConfig) -> None: # side-effects in backward, but torch.compile has no easy way to reapply # python mutations in the backward. Setting this flag accepts this eager # and compile divergence by skipping reapplication of side effects. - torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = ( - True # pyrefly: ignore [bad-assignment] - ) + torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True backend = _maybe_regional_inductor_backend(model, compile_config.backend) @@ -64,26 +63,23 @@ def _maybe_regional_inductor_backend(model: nn.Module, backend: str) -> str | Ca ``regional_inductor`` lowers just the regions annotated with ``compile_with_inductor`` (see ``FlexAttention.forward``) to inductor while the rest stays in aot_eager. - Only applied for ``aot_eager`` on models that actually use FlexAttention, so - dense/non-flex aot_eager paths are left untouched. Other non-inductor backends - can't be scooped here and raise rather than silently degrading. + Only applied for ``aot_eager`` on models that actually use FlexAttention. + The ``inductor`` and ``neuron`` backends handle FlexAttention directly. + Other backends raise when the model uses FlexAttention. """ from torchtitan.models.common.attention import FlexAttention uses_flex = any(isinstance(m, FlexAttention) for m in model.modules()) - # Non-flex models never need the scoop; the default inductor backend already - # lowers the flex region directly. Both are left on the unmodified backend. - if not uses_flex or backend == "inductor": + # Non-flex models do not need regional compilation; inductor and neuron + # already handle the FlexAttention region directly. + if not uses_flex or backend in ("inductor", "neuron"): return backend - # FlexAttention only has an inductor lowering. Under a non-inductor backend - # other than aot_eager it would decompose to eager aten ops (no Triton - # kernel), which we can't transparently scoop here -- fail loudly. if backend != "aot_eager": raise ValueError( - f"Model uses FlexAttention but compile backend {backend!r} is neither " - f"'inductor' nor 'aot_eager'; the flex region would decompose to eager " - f"aten ops (no Triton kernel). Use 'inductor' or 'aot_eager'." + f"Model uses FlexAttention but compile backend {backend!r} is not " + "supported. Use 'inductor', 'neuron', or 'aot_eager' with regional " + "Inductor compilation." ) from torch._dynamo.backends.common import aot_autograd diff --git a/torchtitan/distributed/context_parallel.py b/torchtitan/distributed/context_parallel.py index bcda8d0deba..50ab02272a7 100644 --- a/torchtitan/distributed/context_parallel.py +++ b/torchtitan/distributed/context_parallel.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from collections.abc import Sequence from typing import Any, cast @@ -29,9 +30,39 @@ ScaledDotProductAttention, VarlenAttention, ) +from torchtitan.models.common.rope import RoPE from torchtitan.tools.logging import logger +def apply_cp_to_rope( + model: nn.Module, + cp_mesh: DeviceMesh, + *, + load_balancer_type: str | None, +) -> None: + """Configure prefix RoPE caches to follow the input's CP sharding.""" + if cp_mesh.ndim != 1 or cp_mesh.mesh_dim_names != ("cp",): + raise ValueError("RoPE context parallelism requires a 1D mesh named 'cp'") + if load_balancer_type not in (None, "headtail", "ptrr"): + raise ValueError( + f"Invalid context-parallel load balancer: {load_balancer_type}" + ) + for module in model.modules(): + if isinstance(module, RoPE): + if ( + cp_mesh.size() > 1 + and load_balancer_type == "ptrr" + and module.uses_prefix_cache_for_positions + ): + raise ValueError( + "Prefix RoPE caches require " + "context_parallel_load_balancer=None or 'headtail'" + ) + module._cp_rank = cp_mesh.get_local_rank() + module._cp_size = cp_mesh.size() + module._cp_load_balancer = load_balancer_type + + def apply_cp_to_forward( attention_modules: Sequence[nn.Module], cp_mesh: DeviceMesh, @@ -122,15 +153,14 @@ def prepare_context_parallel_input( """ Shard inputs, labels, positions, and attention masks for Context Parallel. - The caller must provide ``extra_kwargs["positions"]`` before calling this - function. Position resolution (per-document vs sequential) is handled - upstream in ``post_dataloading_process``. + Preserve per-document positions when provided. For loaders that omit + positions, create global sequential positions before sharding. Args: inputs: Input tensor of shape [batch_size, seq_len] labels: Label tensor of shape [batch_size, seq_len] - extra_kwargs: Dictionary containing 'positions' (required) and - optionally 'attention_masks' to be sharded. + extra_kwargs: Dictionary optionally containing 'positions' and + 'attention_masks' to be sharded. cp_mesh: Device mesh for context parallel dimension device: Device for the tensors load_balancer_type: Type of load balancer to use for sharding. @@ -144,7 +174,11 @@ def prepare_context_parallel_input( sharded 'attention_masks' """ attention_masks = extra_kwargs.get("attention_masks", None) - positions = extra_kwargs["positions"] + positions = extra_kwargs.get("positions") + if positions is None: + positions = torch.arange(inputs.shape[1], device=device).expand( + inputs.shape[0], -1 + ) (inputs, labels, positions), attention_masks = cp_shard( cp_mesh, (inputs, labels, positions), diff --git a/torchtitan/distributed/neuron_ep/README.md b/torchtitan/distributed/neuron_ep/README.md new file mode 100644 index 00000000000..a7080048c6d --- /dev/null +++ b/torchtitan/distributed/neuron_ep/README.md @@ -0,0 +1,15 @@ +# Neuron Expert Parallelism + +This package integrates the NKI Library (`nkilib`) `permute_a2av` and +`unpermute_a2av` kernels with TorchTitan's token dispatcher and autograd. + +The packed dispatch and combine kernels are imported from the NKI Library. +Their autograd wrappers and token-dispatcher integration live here. + +## Scope and known limitation + +The `neuron_ep` dispatcher currently supports the published packed-kernel contract: +Trn3 LNC=2, EP >= 8, and one local expert per EP rank. Its receive capacity is +shape-static and grouped-mm aligned. Standard and local dispatch do not add +Neuron-specific grouped-mm padding in core. A follow-up is required for +multi-local-expert regrouping. diff --git a/torchtitan/distributed/neuron_ep/__init__.py b/torchtitan/distributed/neuron_ep/__init__.py new file mode 100644 index 00000000000..51a1175b496 --- /dev/null +++ b/torchtitan/distributed/neuron_ep/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +from torchtitan.distributed.neuron_ep.api import ( + build_ep_replica_group_ranks, + build_packed_source_indices, + gmm_align, + permute_a2av_dispatch_compact, + unpermute_a2av_combine_packed, +) + +__all__ = [ + "build_ep_replica_group_ranks", + "build_packed_source_indices", + "gmm_align", + "permute_a2av_dispatch_compact", + "unpermute_a2av_combine_packed", +] diff --git a/torchtitan/distributed/neuron_ep/api.py b/torchtitan/distributed/neuron_ep/api.py new file mode 100644 index 00000000000..ac430857d1a --- /dev/null +++ b/torchtitan/distributed/neuron_ep/api.py @@ -0,0 +1,284 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron NKI all-to-all-v primitives for the ``neuron_ep`` token dispatcher. + +Wraps the ``permute_a2av`` (dispatch) / ``unpermute_a2av`` (combine) NKI kernels +from ``nkilib.experimental.collectives.a2av_train`` into autograd-aware ops, and +builds the ``nki.collectives.ReplicaGroup`` from an EP ``DeviceMesh``. + +The packed kernels are mathematical transposes: the adjoint of +``permute_a2av`` is ``unpermute_a2av`` and vice-versa. We wrap ONLY the kernel +calls in an ``autograd.Function`` so the surrounding score-reduction stays in +the autograd-tracked graph and the kernel-boundary backward is just the sibling +kernel. + +NKI / nkilib / torch_neuronx imports are deferred to call time so this module +imports on a CPU-only host (the ``NeuronEPTokenDispatcher`` class is importable +everywhere; only running dispatch/combine requires the Neuron toolchain). + +Tensor shape suffixes used here: + M routed rows on the local rank (== N = T*top_k for neuron_ep) + H hidden dimension + E EP degree (number of expert-parallel ranks) + e local experts per rank (num_experts // E) +""" + +from __future__ import annotations + +import torch + + +def gmm_align() -> int: + """Row alignment the Neuron grouped-mm requires for its token dim and offsets. + + Must match the value ``validate_offs`` checks in + ``torch_neuronx.python_ops.grouped_mm`` (default 128, overridable via + ``TORCH_NEURONX_ATEN_GMM_ALIGN``). Falls back to the env/default off-device so + the dispatcher pads consistently even where torch_neuronx is unavailable. + """ + try: + from torch_neuronx.utils import get_gmm_align + + return get_gmm_align() + except ImportError: + import os + + return int(os.environ.get("TORCH_NEURONX_ATEN_GMM_ALIGN", "128")) + + +_LNC: int | None = None + + +def _resolve_lnc() -> int: + """Return the cached LNC, resolving it once before graph capture.""" + global _LNC + if _LNC is None: + from torch_neuronx.utils import get_logical_neuron_cores + + _LNC = int(get_logical_neuron_cores()) + return _LNC + + +def build_ep_replica_group_ranks(ep_mesh) -> list[list[int]]: + """Return the full SPMD replica-group rank lists for an EP ``DeviceMesh``. + + ``nki.collectives.ReplicaGroup`` is constructed from ``list[list[int]]`` -- + one rank list per replica group. Under SPMD every rank must see the SAME + full set of groups, so we read them from the Neuron mesh registry (populated + by the DeviceMesh init hook at ``torch_neuronx`` import) rather than from the + local process group, which only knows the current rank's group. + """ + from torch_neuronx.distributed import mesh_registry + + _resolve_lnc() + pg_name = ep_mesh.get_group().group_name + # get_all_replica_groups falls back to the current rank's group only when the + # PG is not in the registry -- wrong for multi-node SPMD. Require the entry. + if pg_name not in mesh_registry._MESH_REGISTRY: + raise ValueError( + f"EP process group '{pg_name}' is not in the Neuron mesh registry. " + "neuron_ep needs the full SPMD replica groups. Ensure the EP mesh was " + "created via DeviceMesh so the registry hook captured it, and that " + "TORCH_NEURONX_SPMD_DISABLE is unset." + ) + ep_ranks = mesh_registry.get_all_replica_groups(pg_name) + _validate_packed_ep_ranks(ep_ranks) + return ep_ranks + + +def _validate_packed_ep_ranks(ep_ranks: list[list[int]]) -> None: + """Validate the Trn3 device-group topology required by packed A2AV.""" + if not ep_ranks: + raise ValueError("packed A2AV requires at least one EP replica group.") + for ranks in ep_ranks: + ep_size = len(ranks) + if ep_size < 8 or ep_size % 4 != 0: + raise ValueError( + "packed A2AV requires EP >= 8 in complete four-rank Trn3 " + f"device groups. Got replica group of size {ep_size}." + ) + expected = list(range(ranks[0], ranks[0] + ep_size)) + if ranks != expected or ranks[0] % 4 != 0: + raise ValueError( + "packed A2AV requires sequential, four-rank-device-aligned " + f"replica groups. Got {ranks}." + ) + + +def _packed_lnc2_grid(logical_neuron_cores: int) -> tuple[int]: + if logical_neuron_cores != 2: + raise RuntimeError( + "packed A2AV requires NEURON_LOGICAL_NC_CONFIG=2; " + f"got {logical_neuron_cores} logical NeuronCore(s)." + ) + return (logical_neuron_cores,) + + +def _packed_permute_a2av_kernel( + hidden_states_NH, + packed_source_indices_N, + send_counts_1E, + recv_capacity, + ep_ranks, +): + """Call packed dispatch, returning ``[C,H]`` data and A2AV metadata.""" + from nki.collectives import ReplicaGroup + from nkilib.experimental.collectives.a2av_train import permute_a2av + from torch_neuronx.nki_hop import wrap_nki + + grid = _packed_lnc2_grid(_resolve_lnc()) + return wrap_nki(permute_a2av)[grid]( + hidden_states=hidden_states_NH, + send_counts=send_counts_1E, + recv_capacity=int(recv_capacity), + replica_group=ReplicaGroup(ep_ranks), + packed_source_indices=packed_source_indices_N, + ) + + +def build_packed_source_indices(expert_ids_N, send_counts_1E): + """Build stable peer segments using tile-local NKI prefix scans.""" + from nkilib.experimental.collectives.a2av_train.build_packed_routing import ( + build_packed_routing, + ) + from torch_neuronx.nki_hop import wrap_nki + + grid = _packed_lnc2_grid(_resolve_lnc()) + return wrap_nki(build_packed_routing)[grid]( + expert_ids_N.to(torch.int32), send_counts_1E + ) + + +def _packed_unpermute_a2av_kernel( + output_CH, + packed_source_indices_N, + send_counts_1E, + recv_counts_1E, + ep_ranks, +): + """Call packed combine, returning rows in original ``[N,H]`` order.""" + from nki.collectives import ReplicaGroup + from nkilib.experimental.collectives.a2av_train import unpermute_a2av + from torch_neuronx.nki_hop import wrap_nki + + grid = _packed_lnc2_grid(_resolve_lnc()) + return wrap_nki(unpermute_a2av)[grid]( + output=output_CH, + send_counts=send_counts_1E, + recv_counts=recv_counts_1E, + replica_group=ReplicaGroup(ep_ranks), + packed_source_indices=packed_source_indices_N, + ) + + +class _PackedPermuteA2AV(torch.autograd.Function): + """Capacity-bounded packed dispatch with packed combine as its adjoint.""" + + @staticmethod + def forward( + ctx, + hidden_states_NH, + packed_source_indices_N, + send_counts_1E, + recv_capacity, + ep_ranks, + ): + recv_data_CH, metadata_4E = _packed_permute_a2av_kernel( + hidden_states_NH, + packed_source_indices_N, + send_counts_1E, + recv_capacity, + ep_ranks, + ) + h = hidden_states_NH.shape[1] + # Metadata is uint32 elements. PyTorch does not implement floor division + # for uint32, and narrowing first can wrap above INT32_MAX. + recv_counts_1E = (metadata_4E[2:3, :].to(torch.int64) // h).to(torch.int32) + ctx.save_for_backward(packed_source_indices_N, send_counts_1E, recv_counts_1E) + ctx.ep_ranks = ep_ranks + return recv_data_CH, recv_counts_1E + + @staticmethod + def backward(ctx, grad_recv_data_CH, _grad_recv_counts): + packed_source_indices_N, send_counts_1E, recv_counts_1E = ctx.saved_tensors + grad_hidden_NH = _packed_unpermute_a2av_kernel( + grad_recv_data_CH.contiguous(), + packed_source_indices_N, + send_counts_1E, + recv_counts_1E, + ctx.ep_ranks, + ) + return grad_hidden_NH, None, None, None, None + + +class _PackedUnpermuteA2AV(torch.autograd.Function): + """Packed combine with capacity-bounded packed dispatch as its adjoint.""" + + @staticmethod + def forward( + ctx, + output_CH, + packed_source_indices_N, + send_counts_1E, + recv_counts_1E, + ep_ranks, + ): + ctx.save_for_backward(packed_source_indices_N, send_counts_1E) + ctx.recv_capacity = output_CH.shape[0] + ctx.ep_ranks = ep_ranks + return _packed_unpermute_a2av_kernel( + output_CH, + packed_source_indices_N, + send_counts_1E, + recv_counts_1E, + ep_ranks, + ) + + @staticmethod + def backward(ctx, grad_combined_NH): + packed_source_indices_N, send_counts_1E = ctx.saved_tensors + grad_output_CH, _metadata = _packed_permute_a2av_kernel( + grad_combined_NH.contiguous(), + packed_source_indices_N, + send_counts_1E, + ctx.recv_capacity, + ctx.ep_ranks, + ) + return grad_output_CH, None, None, None, None + + +def permute_a2av_dispatch_compact( + hidden_states_NH, + packed_source_indices_N, + send_counts_1E, + recv_capacity, + ep_ranks, +): + """Packed dispatch: ``[N,H]`` -> ``[C,H]`` plus source receive counts.""" + return _PackedPermuteA2AV.apply( + hidden_states_NH, + packed_source_indices_N, + send_counts_1E, + recv_capacity, + ep_ranks, + ) + + +def unpermute_a2av_combine_packed( + output_CH, + packed_source_indices_N, + send_counts_1E, + recv_counts_1E, + ep_ranks, +): + """Packed combine: capacity buffer ``[C,H]`` -> original rows ``[N,H]``.""" + return _PackedUnpermuteA2AV.apply( + output_CH, + packed_source_indices_N, + send_counts_1E, + recv_counts_1E, + ep_ranks, + ) diff --git a/torchtitan/distributed/parallel_dims.py b/torchtitan/distributed/parallel_dims.py index 34d7b580092..773f6b53d81 100644 --- a/torchtitan/distributed/parallel_dims.py +++ b/torchtitan/distributed/parallel_dims.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from __future__ import annotations @@ -230,6 +231,7 @@ def build_mesh(self) -> DeviceMesh: ["pp", "batch", "cp", "tp"] # dataloading_mesh ["pp", "dp_replicate", "dp_shard", "cp", "tp"] # full_dtensor dense_mesh ["pp", "dp_replicate", "fsdp", "tp"] # legacy dense_mesh + ["pp", "fsdp", "dp_replicate", "tp"] # legacy dense_mesh_fsdp_first ["pp", "dp", "cp", "tp"] # spmd_types dense_mesh ["pp", "dp_replicate", "efsdp", "ep"] # sparse_mesh @@ -283,6 +285,7 @@ def unflatten_mesh( ) loss_mesh = dataloading_mesh["batch", "cp"]._flatten("loss_mesh") spmd_dense_mesh_for_fwdbwd = None + dense_mesh_fsdp_first = None if self.spmd_backend == "full_dtensor": # Under full_dtensor, ``dp_shard`` and ``cp`` cannot be folded # together: activations carry a ``cp`` dimension, so parameters @@ -327,6 +330,19 @@ def unflatten_mesh( ("pp", "dp_replicate", "fsdp", "tp"), (self.pp, self.dp_replicate, fsdp, self.tp), ) + # Same devices as the dense mesh, but with ``fsdp`` as the outer + # axis and ``dp_replicate`` inner. SimpleFSDP (graph_trainer) + # requests this mesh via ``get_mesh(["fsdp", "dp_replicate"])`` so + # that DTensor's redistribute processes the ``fsdp`` reduce-scatter + # before the ``dp_replicate`` all-reduce in the backward pass. That + # ordering shards each gradient before the replicate all-reduce + # instead of all-reducing the full unsharded gradient (which, for a + # TP-sharded LM-head weight, is a large fp32 buffer that overflows HBM). + dense_mesh_fsdp_first = unflatten_mesh( + self._world_mesh, + ("pp", "fsdp", "dp_replicate", "tp"), + (self.pp, fsdp, self.dp_replicate, self.tp), + ) full_sparse_mesh = unflatten_mesh( self._world_mesh, @@ -340,6 +356,8 @@ def unflatten_mesh( "dense": full_dense_mesh_for_fsdp, "sparse": full_sparse_mesh, } + if dense_mesh_fsdp_first is not None: + self._global_meshes["dense_fsdp_first"] = dense_mesh_fsdp_first if spmd_dense_mesh_for_fwdbwd is not None: self._global_meshes["spmd_dense_for_fwdbwd"] = spmd_dense_mesh_for_fwdbwd if self.spmd_backend == "spmd_types" and self.ep > 1: @@ -466,15 +484,31 @@ def get_optional_mesh( if key in self._multi_axis_meshes: return self._multi_axis_meshes[key] + # Only the explicit fsdp-first HSDP request -- both DP axes present with + # ``fsdp`` before ``dp_replicate`` -- selects the dedicated + # ``dense_fsdp_first`` mesh. Every other request resolves against the + # original meshes exactly as before, so no existing consumer's + # resolution depends on ``_global_meshes`` insertion order. + fsdp_first_mesh = self._global_meshes.get("dense_fsdp_first") + wants_fsdp_first = ( + "fsdp" in dims + and "dp_replicate" in dims + and dims.index("fsdp") < dims.index("dp_replicate") + ) candidates = [ - (name, global_mesh) - for name, global_mesh in self._global_meshes.items() + global_mesh + for global_mesh in self._global_meshes.values() if global_mesh.mesh_dim_names is not None and set(dims).issubset(set(global_mesh.mesh_dim_names)) + and (wants_fsdp_first or global_mesh is not fsdp_first_mesh) ] if not candidates: raise ValueError(f"Invalid mesh name combinations {dims}.") - submesh = candidates[0][1][key] + if wants_fsdp_first and fsdp_first_mesh in candidates: + chosen = fsdp_first_mesh + else: + chosen = candidates[0] + submesh = chosen[key] self._multi_axis_meshes[key] = submesh return submesh diff --git a/torchtitan/distributed/utils.py b/torchtitan/distributed/utils.py index 6a2018e0f48..a4f9fb9effc 100644 --- a/torchtitan/distributed/utils.py +++ b/torchtitan/distributed/utils.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from __future__ import annotations @@ -10,7 +11,8 @@ import math import os from abc import abstractmethod -from collections.abc import Iterable +from collections.abc import Callable, Iterable +from dataclasses import dataclass from datetime import timedelta from typing import Protocol, TYPE_CHECKING @@ -25,8 +27,8 @@ from torch.distributed.tensor import DTensor from torch.distributed.tensor.placement_types import Placement, Shard -from torchtitan.config import CommConfig, DebugConfig -from torchtitan.tools.logging import logger +from torchtitan.config import CommConfig, CompileConfig, DebugConfig +from torchtitan.tools.logging import logger, warn_once from torchtitan.tools.utils import device_module, device_type if TYPE_CHECKING: @@ -549,6 +551,145 @@ def set_pg_timeouts( torch.distributed.set_timeout(timeout, group) +def _compute_grouped_norm_sq( + grads: list[torch.Tensor], + weights: torch.Tensor, + out: torch.Tensor, +) -> None: + out.zero_() + for i, grad in enumerate(grads): + grad_fp32 = grad.to(torch.float32) + out.add_((grad_fp32 * grad_fp32).sum() * weights[i]) + + +def _compute_clip_coef( + local_norm_sq: torch.Tensor, + grads: list[torch.Tensor], + max_norm: float, + total_norm_out: torch.Tensor, +) -> None: + total_norm_out.copy_(local_norm_sq.sqrt()) + coef = torch.clamp(max_norm / (total_norm_out + 1e-6), max=1.0) + for grad in grads: + grad.mul_(coef) + + +def _replication_weight(param: DTensor | torch.Tensor) -> float: + if not isinstance(param, DTensor): + return 1.0 + + replication = 1 + for mesh_axis, placement in enumerate(param.placements): + if placement.is_replicate(): + replication *= param.device_mesh.size(mesh_axis) + return 1.0 / replication + + +@dataclass +class _CompiledGradClipState: + weights: torch.Tensor + local_norm_sq: torch.Tensor + total_norm: torch.Tensor + + +class CompiledGradClip: + """Compiled gradient clipping using the configured torch.compile backend.""" + + def __init__(self, backend: str) -> None: + self._backend = backend + self._norm_sq_compiled: Callable[..., None] | None = None + self._coef_scale_compiled: Callable[..., None] | None = None + self._state_cache: dict[tuple[int, ...], _CompiledGradClipState] = {} + + def _get_state( + self, + params_with_grad: list[torch.Tensor], + device: torch.device, + ) -> _CompiledGradClipState: + parameter_key = tuple(id(param) for param in params_with_grad) + state = self._state_cache.get(parameter_key) + if state is not None: + return state + + weights = torch.tensor( + [_replication_weight(param) for param in params_with_grad], + dtype=torch.float32, + device=device, + ) + state = _CompiledGradClipState( + weights=weights, + local_norm_sq=torch.zeros((), dtype=torch.float32, device=device), + total_norm=torch.zeros((), dtype=torch.float32, device=device), + ) + self._state_cache[parameter_key] = state + return state + + @torch.no_grad() + def __call__( + self, + parameters: list[torch.Tensor], + max_norm: float, + ) -> torch.Tensor: + params_with_grad = [param for param in parameters if param.grad is not None] + local_grads = [] + for param in params_with_grad: + grad = param.grad + assert grad is not None + local_grads.append(grad.to_local() if isinstance(grad, DTensor) else grad) + + if local_grads: + device = local_grads[0].device + elif parameters: + first_parameter = parameters[0] + device = ( + first_parameter.to_local().device + if isinstance(first_parameter, DTensor) + else first_parameter.device + ) + else: + device = torch.device(device_type) + + state = self._get_state(params_with_grad, device) + + if local_grads: + if self._norm_sq_compiled is None: + self._norm_sq_compiled = torch.compile( + _compute_grouped_norm_sq, + backend=self._backend, + dynamic=False, + ) + self._norm_sq_compiled(local_grads, state.weights, state.local_norm_sq) + else: + state.local_norm_sq.zero_() + dist.all_reduce(state.local_norm_sq, op=dist.ReduceOp.SUM) + + if local_grads: + if self._coef_scale_compiled is None: + self._coef_scale_compiled = torch.compile( + _compute_clip_coef, + backend=self._backend, + dynamic=False, + ) + self._coef_scale_compiled( + state.local_norm_sq, + local_grads, + max_norm, + state.total_norm, + ) + else: + state.total_norm.copy_(state.local_norm_sq.sqrt()) + return state.total_norm + + +def build_compiled_grad_clip( + compile_config: CompileConfig, +) -> CompiledGradClip | None: + """Build compiled gradient clipping when selected by the compile config.""" + if compile_config.enable and "grad_clip" in compile_config.components: + return CompiledGradClip(compile_config.backend) + return None + + @torch.no_grad() def clip_grad_norm_( parameters: torch.Tensor | Iterable[torch.Tensor], @@ -558,6 +699,7 @@ def clip_grad_norm_( foreach: bool | None = None, pp_mesh: DeviceMesh | None = None, ep_enabled: bool = False, + compiled_grad_clip: CompiledGradClip | None = None, ) -> torch.Tensor: """ Clip the gradient norm of an iterable of parameters. @@ -582,11 +724,37 @@ def clip_grad_norm_( pp_mesh: Pipeline Parallel device mesh. If not None, will reduce gradient norm across PP stages. ep_dense_params_mesh_ndim: Mesh ndim of the dense params when EP is used. If EP is not used, set it to ``None``. + compiled_grad_clip: optional compiled gradient clipping implementation. Returns: Total norm of the parameter gradients (viewed as a single vector). """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + else: + parameters = list(parameters) + + if compiled_grad_clip is not None: + unsupported_reasons = [] + if pp_mesh is not None: + unsupported_reasons.append("pipeline parallelism is enabled") + if ep_enabled: + unsupported_reasons.append("expert parallelism is enabled") + if norm_type != 2.0: + unsupported_reasons.append(f"norm_type is {norm_type}, not 2.0") + if error_if_nonfinite: + unsupported_reasons.append("error_if_nonfinite is enabled") + + if not unsupported_reasons: + return compiled_grad_clip(parameters, max_norm) + + warn_once( + logger, + "Compiled gradient clipping was requested but cannot be applied because " + f"{', '.join(unsupported_reasons)}; falling back to eager clipping.", + ) + if ep_enabled: return _clip_grad_norm_with_ep( parameters, @@ -597,11 +765,6 @@ def clip_grad_norm_( pp_mesh, ) - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - else: - # prevent generators from being exhausted - parameters = list(parameters) grads = [p.grad for p in parameters if p.grad is not None] total_norm = torch.nn.utils.get_total_norm( grads, norm_type, error_if_nonfinite, foreach diff --git a/torchtitan/experiments/forge/engine.py b/torchtitan/experiments/forge/engine.py index f0a3adf5618..78af445cb02 100644 --- a/torchtitan/experiments/forge/engine.py +++ b/torchtitan/experiments/forge/engine.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import os from collections.abc import Generator @@ -89,6 +90,7 @@ def to_dict(self) -> dict[str, Any]: # runtime utilities device: torch.device gc_handler: utils.GarbageCollection + compiled_grad_clip: dist_utils.CompiledGradClip | None gradient_accumulation_steps: int train_context: Generator[None, None, None] pp_has_first_stage: bool @@ -115,6 +117,7 @@ def __init__(self, config: Config): self.device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}") # Device has to be set before creating TorchFT manager. device_module.set_device(self.device) + self.compiled_grad_clip = dist_utils.build_compiled_grad_clip(config.compile) # init distributed and build meshes dist_utils.init_distributed( diff --git a/torchtitan/experiments/forge/example_train.py b/torchtitan/experiments/forge/example_train.py index 02f05336201..77f0ce819b8 100644 --- a/torchtitan/experiments/forge/example_train.py +++ b/torchtitan/experiments/forge/example_train.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import os import time @@ -306,6 +307,7 @@ def train_step( foreach=True, pp_mesh=parallel_dims.get_optional_mesh("pp"), ep_enabled=parallel_dims.ep_enabled, + compiled_grad_clip=self.compiled_grad_clip, ) self.checkpointer.maybe_wait_for_staging() self.optimizers.step() diff --git a/torchtitan/experiments/graph_trainer/common_utils.py b/torchtitan/experiments/graph_trainer/common_utils.py index 531b675b52d..37ee2c85f34 100644 --- a/torchtitan/experiments/graph_trainer/common_utils.py +++ b/torchtitan/experiments/graph_trainer/common_utils.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import fnmatch import time @@ -43,32 +44,20 @@ def log_timer(label: str): def build_decoder_config_for_backend( config_builder: Callable, attn_backend: str, **builder_kwargs ): - """Build a Decoder model config for ``attn_backend``, allowing test-only SDPA. - - ``SDPA`` is not a valid production language-model backend — ``get_attention_config`` - rejects it because the dataloaders always emit per-document positions and SDPA - cannot consume them (it only has a boolean ``is_causal``). The graph_trainer - tests, however, use SDPA to exercise *backend-agnostic* graph machinery - (precompile-artifact serialization, custom codegen, context parallel, bitwise - determinism) without FlexAttention's ``BlockMask``, which is unpicklable (its - ``mask_mod`` closures are Python code objects), is not a tensor (so it breaks - pipeline-parallel split-backward, which calls ``.requires_grad`` on every stage - input), and overflows the fp32 Triton shared-memory limit on large head dims. - - For SDPA we build the flex config (a valid backend) and swap each layer's - ``inner_attention`` to ``ScaledDotProductAttention.Config()``. Production code - never reaches this path: ``get_attention_config`` still rejects ``sdpa``, so no - model registry can construct an SDPA language model outside these tests. + """Build a Decoder model config for ``attn_backend``. + + SDPA is a first-class backend (``get_attention_config`` constructs it + directly), so this is a plain passthrough. The graph_trainer config + presets select SDPA (the registries themselves default to flex) to + exercise *backend-agnostic* graph machinery (precompile-artifact + serialization, custom codegen, context parallel, bitwise determinism) + without FlexAttention's ``BlockMask``, which is unpicklable (its + ``mask_mod`` closures are Python code objects), is not a tensor (so it + breaks pipeline-parallel split-backward, which calls ``.requires_grad`` on + every stage input), and overflows the fp32 Triton shared-memory limit on + large head dims. """ - if attn_backend != "sdpa": - return config_builder(attn_backend=attn_backend, **builder_kwargs) - - from torchtitan.models.common.attention import ScaledDotProductAttention - - config = config_builder(attn_backend="flex", **builder_kwargs) - for layer in config.layers: - layer.attention.inner_attention = ScaledDotProductAttention.Config() - return config + return config_builder(attn_backend=attn_backend, **builder_kwargs) def _local_stride(tensor: torch.Tensor) -> tuple[int, ...]: @@ -412,15 +401,28 @@ def apply_simple_fsdp( *, parallel_dims: ParallelDims, training: TrainingConfig, + fsdp_first: bool = False, ) -> nn.Module: """Wrap the model (and any MoE experts) with graph_trainer's simple_fsdp. For MoE-enabled models, ``moe.experts`` submodules are separately wrapped on the EDP mesh when expert parallelism is enabled. + + When ``fsdp_first`` is set, the HSDP DP mesh is requested with the ``fsdp`` + axis outer (``["fsdp", "dp_replicate"]``) so SimpleFSDP's backward + reduce-scatters over ``fsdp`` before the ``dp_replicate`` all-reduce, + instead of all-reducing the full unsharded gradient first -- which for a + TP-sharded LM-head weight is a large fp32 buffer that overflows a 24 GiB + core. The default keeps the upstream replicate-first order. + Only the main model opts in; the MoE experts EDP mesh always uses the + default order (its placement is derived by axis name in ``data_parallel``, + so it is correct either way). """ if parallel_dims.dp_replicate_enabled: if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled: - dp_mesh_dim_names = ["dp_replicate", "fsdp"] + dp_mesh_dim_names = ( + ["fsdp", "dp_replicate"] if fsdp_first else ["dp_replicate", "fsdp"] + ) dp_mode = "hybrid_shard" else: dp_mesh_dim_names = ["dp_replicate"] @@ -445,6 +447,11 @@ def apply_simple_fsdp( assert edp_mesh is not None for _, transformer_block in model.layers.items(): + # FullAC replaces each TransformerBlock with a checkpoint wrapper. + # Inspect the wrapped block so its MoE experts still receive EFSDP. + transformer_block = getattr( + transformer_block, "_checkpoint_wrapped_module", transformer_block + ) if not isinstance(transformer_block, TransformerBlock): continue moe = getattr(transformer_block, "moe", None) diff --git a/torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py b/torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py index f620af62008..75fbec37152 100644 --- a/torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py +++ b/torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from torchtitan.config import ParallelismConfig, TrainingConfig from torchtitan.distributed import ParallelDims @@ -81,6 +82,11 @@ def parallelize_deepseekv3( if parallel_dims.tp_enabled: maybe_enable_async_tp(parallelism, compile_config, parallel_dims.get_mesh("tp")) + if compile_config.mode == "jit" and ac_config is not None: + # JIT does not run the AOT memory-policy/remat passes, so honor the + # Trainer AC policy with eager wrappers before SimpleFSDP and compile. + ac_config.build(dump_folder=dump_folder).apply(model) + # Apply simple_fsdp unconditionally. The `fsdp` mesh always exists with a # real backend (see ParallelDims._mesh_exist), even at degree 1, so that # MixedPrecisionPolicy's param_dtype cast still applies in single-GPU runs. diff --git a/torchtitan/experiments/graph_trainer/llama3/__init__.py b/torchtitan/experiments/graph_trainer/llama3/__init__.py index d2036f005ae..3095cacf932 100644 --- a/torchtitan/experiments/graph_trainer/llama3/__init__.py +++ b/torchtitan/experiments/graph_trainer/llama3/__init__.py @@ -1,14 +1,17 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from dataclasses import fields from torchtitan.experiments.graph_trainer.graph_pp.pipeline import graph_pipeline_llm from torchtitan.models.llama3 import llama3_configs from torchtitan.models.llama3.state_dict_adapter import Llama3StateDictAdapter +from torchtitan.models.utils import validate_converter_order +from torchtitan.protocols.model import ModelConfigConverter from torchtitan.protocols.model_spec import ModelSpec from ..common_utils import build_decoder_config_for_backend @@ -29,11 +32,16 @@ def _parallelize_fn(model, *, compile_config, **kwargs): def model_registry( flavor: str, attn_backend: str = "flex", + converters: list[ModelConfigConverter.Config] | None = None, ) -> ModelSpec: base = build_decoder_config_for_backend(llama3_configs[flavor], attn_backend) config = GraphTrainerLlama3Model.Config( **{f.name: getattr(base, f.name) for f in fields(base)} ) + if converters is not None: + validate_converter_order(converters) + for c in converters: + config = c.build().convert(config) return ModelSpec( name="graph_trainer/llama3", flavor=flavor, diff --git a/torchtitan/experiments/graph_trainer/llama3/parallelize.py b/torchtitan/experiments/graph_trainer/llama3/parallelize.py index d9046724f44..b35efd2d874 100644 --- a/torchtitan/experiments/graph_trainer/llama3/parallelize.py +++ b/torchtitan/experiments/graph_trainer/llama3/parallelize.py @@ -1,10 +1,11 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. -from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.config import ParallelismConfig, TrainingConfig from torchtitan.distributed import ParallelDims from torchtitan.distributed.activation_checkpoint import ActivationCheckpointingConfig from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp @@ -14,6 +15,7 @@ apply_simple_fsdp, ) from torchtitan.experiments.graph_trainer.compile import apply_compile +from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig from torchtitan.experiments.graph_trainer.llama3.model import GraphTrainerLlama3Model @@ -33,7 +35,7 @@ def parallelize_llama( parallel_dims: ParallelDims, training: TrainingConfig, parallelism: ParallelismConfig, - compile_config: CompileConfig, + compile_config: GraphTrainerCompileConfig, ac_config: ActivationCheckpointingConfig, dump_folder: str, ): @@ -60,10 +62,20 @@ def parallelize_llama( model.parallelize(parallel_dims) maybe_enable_async_tp(parallelism, compile_config, parallel_dims.get_mesh("tp")) + if compile_config.mode == "jit" and ac_config is not None: + # JIT does not run the AOT memory-policy/remat passes, so honor the + # Trainer AC policy with eager wrappers before SimpleFSDP and compile. + ac_config.build(dump_folder=dump_folder).apply(model) + # Apply simple_fsdp unconditionally. The `fsdp` mesh always exists with a # real backend (see ParallelDims._mesh_exist), even at degree 1, so that # MixedPrecisionPolicy's param_dtype cast still applies in single-GPU runs. - model = apply_simple_fsdp(model, parallel_dims=parallel_dims, training=training) + # fsdp_first: request the fsdp-outer HSDP mesh so SimpleFSDP reduce-scatters + # over fsdp before the dp_replicate all-reduce, avoiding the full unsharded + # LM-head gradient buffer that exceeds a 24 GiB core. + model = apply_simple_fsdp( + model, parallel_dims=parallel_dims, training=training, fsdp_first=True + ) # Apply compilation based on mode model = apply_compile( diff --git a/torchtitan/experiments/graph_trainer/precompile_main.py b/torchtitan/experiments/graph_trainer/precompile_main.py index f6ac75e55c5..fffeeee13d1 100644 --- a/torchtitan/experiments/graph_trainer/precompile_main.py +++ b/torchtitan/experiments/graph_trainer/precompile_main.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. """ Single-process precompile entry point for graph_trainer. @@ -37,6 +38,7 @@ from torchtitan.experiments.graph_trainer.storage import DiskStorageAdapter from torchtitan.models.common.attention import FlexAttention, VarlenAttention from torchtitan.models.common.decoder import Decoder +from torchtitan.models.common.moe import uses_global_batch_lbl from torchtitan.tools import utils from torchtitan.tools.logging import logger @@ -197,6 +199,10 @@ def _precompile_aot_fx_trace( ) from torchtitan.experiments.graph_trainer.trainer import make_fwd_bwd_step + if uses_global_batch_lbl(model): + # The trainer sets the token-count mesh and step normalization after this graph is traced. + raise ValueError("load_balancing='global_aux_loss' does not support precompile") + loss_fn = config.loss.build(compile_config=compile_config) _prepare_loss_for_precompile(model, loss_fn) diff --git a/torchtitan/experiments/graph_trainer/qwen3/parallelize.py b/torchtitan/experiments/graph_trainer/qwen3/parallelize.py index 0c211ea56b8..d51ff8186a6 100644 --- a/torchtitan/experiments/graph_trainer/qwen3/parallelize.py +++ b/torchtitan/experiments/graph_trainer/qwen3/parallelize.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from torchtitan.config import ParallelismConfig, TrainingConfig from torchtitan.distributed import ParallelDims @@ -73,6 +74,11 @@ def parallelize_qwen3( if parallel_dims.tp_enabled: maybe_enable_async_tp(parallelism, compile_config, parallel_dims.get_mesh("tp")) + if compile_config.mode == "jit" and ac_config is not None: + # JIT does not run the AOT memory-policy/remat passes, so honor the + # Trainer AC policy with eager wrappers before SimpleFSDP and compile. + ac_config.build(dump_folder=dump_folder).apply(model) + # Apply simple_fsdp unconditionally. The `fsdp` mesh always exists with a # real backend (see ParallelDims._mesh_exist), even at degree 1, so that # MixedPrecisionPolicy's param_dtype cast still applies in single-GPU runs. diff --git a/torchtitan/experiments/graph_trainer/simple_fsdp.py b/torchtitan/experiments/graph_trainer/simple_fsdp.py index 19516edd647..ca4b1e9fd89 100644 --- a/torchtitan/experiments/graph_trainer/simple_fsdp.py +++ b/torchtitan/experiments/graph_trainer/simple_fsdp.py @@ -1,13 +1,15 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import sys from collections.abc import Generator, Sequence from contextlib import contextmanager from dataclasses import dataclass +from itertools import count import torch import torch.nn as nn @@ -20,8 +22,12 @@ Shard, ) from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor._dtensor_spec import DTensorSpec +from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta from torch.distributed.tensor._redistribute import redistribute_local_tensor +from torch.distributed.tensor._utils import ( + compute_local_shape_and_global_offset, + compute_local_stride, +) from torch.distributed.tensor.placement_types import _StridedShard, Placement from torchtitan.protocols.module import Module @@ -79,19 +85,26 @@ def _distribute_dtensor( assert len(inner_spec.placements) == 1 tensor_placement = (dp_placements[0], inner_spec.placements[0]) elif len(dp_placements) == 2: - assert dp_placements[0].is_replicate() and dp_placements[1].is_shard() - # For HSDP + EP/TP/EP+TP + # HSDP + EP/TP/EP+TP. The two DP placements are one Shard and one + # Replicate; their order follows the DP mesh axis order, which is + # replicate-first for the default ``dense`` mesh and fsdp-first for the + # ``dense_fsdp_first`` mesh that only the llama3 SimpleFSDP path opts + # into. Handle either order by locating the Shard rather than assuming + # a fixed position. + assert ( + sum(p.is_shard() for p in dp_placements) == 1 + and sum(p.is_replicate() for p in dp_placements) == 1 + ), f"HSDP expects one Shard and one Replicate placement, got {dp_placements}" assert len(inner_spec.placements) == 2 or len(inner_spec.placements) == 1 - shard_dim = dp_placements[1].dim + shard_idx = 0 if dp_placements[0].is_shard() else 1 + shard_dim = dp_placements[shard_idx].dim split_factor = inner_spec.num_shards_map[shard_dim] - tensor_placement = ( - dp_placements[0], - ( - _StridedShard(shard_dim, split_factor=split_factor) - if split_factor > 1 - else dp_placements[1] - ), - ) + inner_spec.placements + dp_tensor_placement = list(dp_placements) + if split_factor > 1: + dp_tensor_placement[shard_idx] = _StridedShard( + shard_dim, split_factor=split_factor + ) + tensor_placement = tuple(dp_tensor_placement) + inner_spec.placements else: raise ValueError( f"Unsupported placement {dp_placements} for distributing DTensor {tensor}" @@ -100,16 +113,20 @@ def _distribute_dtensor( # HSDP case needs 2 placements for 2D outer_mesh current_placements = (Replicate(),) * len(dp_placements) target_placements = tuple(dp_placements) + # The DP mesh distributes the TP/EP-local tensor, not the global tensor. + local_meta = TensorMeta( + tensor._local_tensor.shape, tensor._local_tensor.stride(), tensor.dtype + ) current_spec = DTensorSpec( mesh=outer_mesh, placements=current_placements, - tensor_meta=inner_spec.tensor_meta, + tensor_meta=local_meta, ) target_spec = DTensorSpec( mesh=outer_mesh, placements=target_placements, - tensor_meta=inner_spec.tensor_meta, + tensor_meta=local_meta, ) result_tensor = redistribute_local_tensor( tensor._local_tensor, @@ -130,6 +147,10 @@ def _distribute_dtensor( # Cache of (original_class, param_names) -> wrapper class, so all instances # of the same module type share one SimpleFSDP class for torch.compile reuse. _wrap_class_cache: dict[tuple[type, frozenset[str]], type] = {} +# Suffixes each wrapper class name, as upstream does: one module type can need +# several wrappers (nn.Linear with and without a bias), and pickle resolves each +# by name. The id depends only on wrap order, so it matches across processes. +_wrap_class_id = count() def _register_parametrization( @@ -153,7 +174,7 @@ def _register_parametrization( module_cls = _wrap_class_cache[cache_key] else: module_cls = type( - f"SimpleFSDP{module.__class__.__name__}", + f"SimpleFSDP{module.__class__.__name__}_{next(_wrap_class_id)}", (module.__class__,), param_name_to_property, ) @@ -199,10 +220,24 @@ def replicate_compute(self, x: DTensor) -> torch.Tensor: "full_dtensor not implemented for nD parallelisms" ) dp_mesh = self.device_mesh + non_dp_placements = tuple(x._spec.placements[-non_dp_mesh_dims:]) + non_dp_mesh_dim_names = tuple( + x._spec.mesh.mesh_dim_names[-non_dp_mesh_dims:] + ) + non_dp_mesh = x._spec.mesh[non_dp_mesh_dim_names] + # Preserve the true TP/EP-local shape when DP shards are uneven. + dp_shape, _ = compute_local_shape_and_global_offset( + x.shape, non_dp_mesh, non_dp_placements, skip_offset=True + ) + dp_stride = compute_local_stride(x.stride(), dp_shape) # re-wrap 2D DTensor to 1D DTensor on dp_mesh for efficient FSDP all-gather sharded_local_tensor = x.to_local() sharded_dtensor = DTensor.from_local( - sharded_local_tensor, dp_mesh, self.param_sharding + sharded_local_tensor, + dp_mesh, + self.param_sharding, + shape=torch.Size(dp_shape), + stride=dp_stride, ) # the actual FSDP's fwd all-gather & bwd reduce-scatter @@ -219,14 +254,12 @@ def replicate_compute(self, x: DTensor) -> torch.Tensor: grad_placements=self.grad_placements ) - non_dp_placements = tuple(x._spec.placements[-non_dp_mesh_dims:]) - non_dp_mesh_dim_names = tuple( - x._spec.mesh.mesh_dim_names[-non_dp_mesh_dims:] - ) - non_dp_mesh = x._spec.mesh[non_dp_mesh_dim_names] - output = DTensor.from_local( - replicated_local_tensor, non_dp_mesh, non_dp_placements + replicated_local_tensor, + non_dp_mesh, + non_dp_placements, + shape=x.shape, + stride=x.stride(), ) elif non_dp_mesh_dims == 0: output = x.redistribute( @@ -272,11 +305,24 @@ def data_parallel( elif mode == "fully_shard": param_sharding = (Shard(shard_dim),) elif mode == "hybrid_shard": - # replicate inter-host, fully shard intra-host - param_sharding = (Replicate(), Shard(shard_dim)) assert ( device_mesh.ndim == 2 ), "hybrid sharded data parallel requires 2D DeviceMesh" + # Place Replicate on the dp_replicate axis and Shard on the shard axis + # (fsdp/efsdp) by name, so the placement is correct regardless of the DP + # mesh axis order: replicate-first for the default ``dense`` mesh, and + # fsdp-first for the ``dense_fsdp_first`` mesh that only the llama3 + # SimpleFSDP path opts into. Deriving from names avoids coupling the + # placement to a fixed axis order (which would silently invert sharding + # for a differently-ordered mesh, e.g. the MoE experts EDP mesh). + mesh_axis_names = device_mesh.mesh_dim_names + assert ( + mesh_axis_names is not None and "dp_replicate" in mesh_axis_names + ), f"hybrid_shard expects a named 2D DP mesh with a dp_replicate axis, got {mesh_axis_names}" + param_sharding = tuple( + Replicate() if name == "dp_replicate" else Shard(shard_dim) + for name in mesh_axis_names + ) else: raise ValueError(f"Unsupported mode {mode}") @@ -297,7 +343,8 @@ def data_parallel( mod.register_parameter( p_name, nn.Parameter( - distribute_tensor_func(p, device_mesh, param_sharding) + distribute_tensor_func(p, device_mesh, param_sharding), + requires_grad=p.requires_grad, ), ) diff --git a/torchtitan/experiments/graph_trainer/tests/test_numerics.py b/torchtitan/experiments/graph_trainer/tests/test_numerics.py index 7cd6ab5e6c1..3670c71d7bb 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_numerics.py +++ b/torchtitan/experiments/graph_trainer/tests/test_numerics.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import copy import importlib.util @@ -14,13 +15,18 @@ import unittest from collections.abc import Iterator from contextlib import contextmanager +from unittest.mock import patch import torch +import torch.distributed as dist from torch.distributed._composable.fsdp import fully_shard +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import distribute_tensor, Shard from torch.testing._internal.common_fsdp import FSDPTest from torchtitan.components.loss import cross_entropy_loss from torchtitan.distributed import ParallelDims +from torchtitan.experiments.graph_trainer import simple_fsdp from torchtitan.experiments.graph_trainer.simple_fsdp import data_parallel @@ -667,6 +673,71 @@ def test_deepseek_v3_aot_fx_trace_autoparallel_vs_eager(self): class TestSimpleFSDP(FSDPTest): + def test_composed_sharding_metadata(self): + if self.world_size < 4 or self.world_size % 4: + self.skipTest("Requires a multiple of four devices") + device_type = "cuda" if dist.get_backend() == "nccl" else "cpu" + device = ( + torch.device(device_type, self.rank) if device_type == "cuda" else "cpu" + ) + mesh = init_device_mesh( + device_type, + (self.world_size // 4, 2, 2), + mesh_dim_names=("replica", "fsdp", "tp"), + ) + dp_mesh, tp_mesh = mesh["fsdp"], mesh["tp"] + original_redistribute = simple_fsdp.redistribute_local_tensor + + def check_local_metadata(local, *, current_spec, target_spec): + for spec in (current_spec, target_spec): + self.assertEqual(spec.shape, local.shape) + self.assertEqual(spec.tensor_meta.stride, local.stride()) + self.assertEqual(spec.tensor_meta.dtype, local.dtype) + return original_redistribute( + local, current_spec=current_spec, target_spec=target_spec + ) + + # Cover uneven DP, uneven TP, an empty DP shard, and an even control. + for num_rows in (10, 7, 2, 12): + with self.subTest(num_rows=num_rows): + full = torch.arange( + num_rows * 8, dtype=torch.float64, device=device + ).reshape(num_rows, 8) + tp = distribute_tensor(full, tp_mesh, (Shard(0),)) + expected_tp = full.chunk(2, dim=0)[tp_mesh.get_local_rank()] + chunks = expected_tp.chunk(2, dim=0) + dp_rank = dp_mesh.get_local_rank() + expected_shard = ( + chunks[dp_rank] + if dp_rank < len(chunks) + else expected_tp.narrow(0, 0, 0) + ) + with patch.object( + simple_fsdp, "redistribute_local_tensor", check_local_metadata + ): + parameter = torch.nn.Parameter( + simple_fsdp._distribute_dtensor(tp, dp_mesh, (Shard(0),)) + ) + torch.testing.assert_close( + parameter.to_local(), expected_shard, rtol=0, atol=0 + ) + replicate = simple_fsdp.ReplicateComputation( + dp_mesh, (Shard(0),), "fully_shard", None + ) + output = replicate(parameter) + self.assertEqual(output.shape, full.shape) + self.assertEqual(output.stride(), full.stride()) + torch.testing.assert_close( + output.to_local(), expected_tp, rtol=0, atol=0 + ) + output.to_local().square().sum().backward() + torch.testing.assert_close( + parameter.grad.to_local(), + expected_shard * (2 * dp_mesh.size()), + rtol=0, + atol=0, + ) + def init_test(self): self.optimizer = torch.optim.Adam self.loss_fn = cross_entropy_loss diff --git a/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py b/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py index c41d063a218..417a326bfb1 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py +++ b/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py @@ -1,19 +1,210 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. +import pickle import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch import torch import torch.distributed as dist import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh from torchtitan.config.configs import TrainingConfig from torchtitan.distributed import ParallelDims +from torchtitan.distributed.activation_checkpoint import FullAC +from torchtitan.experiments.graph_trainer import simple_fsdp from torchtitan.experiments.graph_trainer.common_utils import apply_simple_fsdp +from torchtitan.experiments.graph_trainer.simple_fsdp import data_parallel +from torchtitan.models.common.decoder import Decoder, TransformerBlock + + +class TestGraphTrainerQwen3JitActivationCheckpoint(unittest.TestCase): + def test_qwen3_jit_full_ac_wraps_layers(self): + from torchtitan.config import ParallelismConfig, TrainingConfig + from torchtitan.experiments.graph_trainer.configs import ( + GraphTrainerCompileConfig, + ) + from torchtitan.experiments.graph_trainer.qwen3 import parallelize + + class FakeParallelDims: + seq_len_divisor = 1 + cp_enabled = False + tp_enabled = False + ep_enabled = False + + class TinyQwen3LikeModel(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([nn.Linear(2, 2)]) + + model = TinyQwen3LikeModel() + original_layer = model.layers[0] + with ( + patch.object(parallelize, "annotate_qwen3"), + patch.object( + parallelize, "apply_simple_fsdp", lambda module, **kwargs: module + ), + patch.object(parallelize, "apply_compile", lambda module, **kwargs: module), + patch.object(parallelize, "maybe_apply_ep_overlap_eager_chunking"), + ): + result = parallelize.parallelize_qwen3( + model, + parallel_dims=FakeParallelDims(), + training=TrainingConfig(seq_len=1), + parallelism=ParallelismConfig(), + compile_config=GraphTrainerCompileConfig(mode="jit"), + ac_config=FullAC.Config(), + dump_folder="dump", + ) + + self.assertIs(result, model) + self.assertIsNot(model.layers[0], original_layer) + self.assertTrue(hasattr(model.layers[0], "_checkpoint_wrapped_module")) + self.assertIs(model.layers[0]._checkpoint_wrapped_module, original_layer) + + +class TestGraphTrainerDeepSeekV3JitActivationCheckpoint(unittest.TestCase): + def test_deepseek_v3_jit_full_ac_wraps_layers(self): + from torchtitan.config import ParallelismConfig, TrainingConfig + from torchtitan.experiments.graph_trainer.configs import ( + GraphTrainerCompileConfig, + ) + from torchtitan.experiments.graph_trainer.deepseek_v3 import parallelize + + class FakeParallelDims: + seq_len_divisor = 1 + cp_enabled = False + tp_enabled = False + ep_enabled = False + + class TinyDeepSeekV3LikeModel(nn.Module): + class Config: + layers = [] + + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([nn.Linear(2, 2)]) + self.config = self.Config() + + model = TinyDeepSeekV3LikeModel() + original_layer = model.layers[0] + + def assert_ac_applied_before_simple_fsdp(module, **_kwargs): + self.assertTrue(hasattr(module.layers[0], "_checkpoint_wrapped_module")) + return module + + with ( + patch.object(parallelize, "annotate_deepseekv3"), + patch.object( + parallelize, + "apply_simple_fsdp", + side_effect=assert_ac_applied_before_simple_fsdp, + ), + patch.object(parallelize, "apply_compile", lambda module, **kwargs: module), + patch.object(parallelize, "maybe_apply_ep_overlap_eager_chunking"), + ): + result = parallelize.parallelize_deepseekv3( + model, + parallel_dims=FakeParallelDims(), + training=TrainingConfig(seq_len=1), + parallelism=ParallelismConfig(), + compile_config=GraphTrainerCompileConfig(mode="jit"), + ac_config=FullAC.Config(), + dump_folder="dump", + ) + + self.assertIs(result, model) + self.assertIsNot(model.layers[0], original_layer) + self.assertTrue(hasattr(model.layers[0], "_checkpoint_wrapped_module")) + self.assertIs(model.layers[0]._checkpoint_wrapped_module, original_layer) + + +class TestGraphTrainerLlama3JitActivationCheckpoint(unittest.TestCase): + def test_llama3_jit_full_ac_wraps_layers(self): + from torchtitan.config import ParallelismConfig, TrainingConfig + from torchtitan.experiments.graph_trainer.configs import ( + GraphTrainerCompileConfig, + ) + from torchtitan.experiments.graph_trainer.llama3 import parallelize + + class FakeParallelDims: + seq_len_divisor = 1 + cp_enabled = False + tp_enabled = False + + class TinyLlama3LikeModel(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([nn.Linear(2, 2)]) + + model = TinyLlama3LikeModel() + original_layer = model.layers[0] + + def assert_ac_applied_before_simple_fsdp(module, **_kwargs): + self.assertTrue(hasattr(module.layers[0], "_checkpoint_wrapped_module")) + return module + + with ( + patch.object(parallelize, "annotate_llama"), + patch.object( + parallelize, + "apply_simple_fsdp", + side_effect=assert_ac_applied_before_simple_fsdp, + ), + patch.object(parallelize, "apply_compile", lambda module, **kwargs: module), + ): + result = parallelize.parallelize_llama( + model, + parallel_dims=FakeParallelDims(), + training=TrainingConfig(seq_len=1), + parallelism=ParallelismConfig(), + compile_config=GraphTrainerCompileConfig(mode="jit"), + ac_config=FullAC.Config(), + dump_folder="dump", + ) + + self.assertIs(result, model) + self.assertIsNot(model.layers[0], original_layer) + self.assertTrue(hasattr(model.layers[0], "_checkpoint_wrapped_module")) + self.assertIs(model.layers[0]._checkpoint_wrapped_module, original_layer) + + +class _Proj(nn.Module): + def __init__(self, bias: bool): + super().__init__() + self.weight = nn.Parameter(torch.ones(2, 2)) + self.bias = nn.Parameter(torch.zeros(2)) if bias else None + + +class TestRegisterParametrizationWrapperNames(unittest.TestCase): + def tearDown(self): + for key in [k for k in simple_fsdp._wrap_class_cache if k[0] is _Proj]: + cls = simple_fsdp._wrap_class_cache.pop(key) + vars(simple_fsdp).pop(cls.__name__, None) + + @staticmethod + def _wrap(module): + names = [n for n, p in module.named_parameters(recurse=False)] + simple_fsdp._register_parametrization(module, names, nn.Identity()) + return type(module) + + def test_wrappers_of_one_type_get_unique_picklable_names(self): + unbiased_cls = self._wrap(_Proj(bias=False)) + biased_cls = self._wrap(_Proj(bias=True)) + + self.assertNotEqual(unbiased_cls.__name__, biased_cls.__name__) + for cls in (unbiased_cls, biased_cls): + self.assertRegex(cls.__name__, r"^SimpleFSDP_Proj_\d+$") + # pickle saves a class by name; each name must resolve to its own class. + self.assertIs(pickle.loads(pickle.dumps(cls)), cls) + + def test_modules_with_same_params_share_one_wrapper(self): + self.assertIs(self._wrap(_Proj(bias=True)), self._wrap(_Proj(bias=True))) class TestApplySimpleFSDPSingleRank(unittest.TestCase): @@ -32,6 +223,30 @@ def tearDown(self): if dist.is_initialized(): dist.destroy_process_group() + def test_frozen_parameter_remains_frozen(self): + # SimpleFSDP caches wrapper classes; keep this no-policy fixture isolated. + class FrozenLinear(nn.Linear): + pass + + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("fsdp",)) + model = FrozenLinear(8, 8) + model.weight.requires_grad_(False) + data_parallel(model, mesh, "fully_shard") + weight = model._parameters["weight"] + bias = model._parameters["bias"] + self.assertFalse(weight.requires_grad) + self.assertTrue(bias.requires_grad) + original_weight = weight.to_local().detach().clone() + original_bias = bias.to_local().detach().clone() + optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) + + model(torch.ones(2, 8)).sum().backward() + self.assertIsNone(weight.grad) + self.assertIsNotNone(bias.grad) + optimizer.step() + torch.testing.assert_close(weight.to_local(), original_weight, rtol=0, atol=0) + self.assertFalse(torch.equal(bias.to_local(), original_bias)) + @patch("torchtitan.distributed.parallel_dims.device_type", "cpu") def test_param_cast_to_bf16_at_ngpu_1(self): """With ``mixed_precision_param=bfloat16``, the parametrized weight must @@ -75,5 +290,61 @@ def test_param_cast_to_bf16_at_ngpu_1(self): self.assertEqual(y.dtype, torch.bfloat16) +class TestApplySimpleFSDPExpertDiscovery(unittest.TestCase): + @patch("torchtitan.experiments.graph_trainer.common_utils.data_parallel") + def test_full_ac_wrapped_block_shards_experts_on_edp_mesh(self, data_parallel): + data_parallel.side_effect = lambda module, *_args, **_kwargs: module + + experts = nn.Module() + experts.num_experts = 64 + moe = nn.Module() + moe.experts = experts + block = TransformerBlock() + block.moe = moe + + model = Decoder.__new__(Decoder) + nn.Module.__init__(model) + model.layers = nn.ModuleDict( + {"0": FullAC.Config().build()._wrap_block(block, base_fqn="layers.0")} + ) + + efsdp_mesh = MagicMock() + efsdp_mesh.size.return_value = 1 + parallel_dims = MagicMock() + parallel_dims.dp_replicate_enabled = False + parallel_dims.dp_shard_enabled = True + parallel_dims.cp_enabled = False + parallel_dims.ep_enabled = True + parallel_dims.ep = 64 + parallel_dims.get_mesh.return_value = MagicMock() + parallel_dims.get_optional_mesh.return_value = {"efsdp": efsdp_mesh} + + training = TrainingConfig( + mixed_precision_param="bfloat16", + mixed_precision_reduce="float32", + ) + + result = apply_simple_fsdp( + model, + parallel_dims=parallel_dims, + training=training, + ) + + self.assertIs(result, model) + self.assertEqual(data_parallel.call_count, 2) + experts_call, model_call = data_parallel.call_args_list + self.assertIs(experts_call.args[0], experts) + self.assertEqual(experts_call.args[1], {"efsdp": efsdp_mesh}) + self.assertEqual(experts_call.args[2], "fully_shard") + self.assertEqual(experts_call.kwargs["shard_dim"], 0) + self.assertIs(model_call.args[0], model) + self.assertIs(model_call.args[1], parallel_dims.get_mesh.return_value) + self.assertEqual(model_call.args[2], "fully_shard") + self.assertIs( + experts_call.kwargs["mp_policy"], + model_call.kwargs["mp_policy"], + ) + + if __name__ == "__main__": unittest.main() diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py index 52588b25f08..3dbcb9fe5fe 100644 --- a/torchtitan/experiments/graph_trainer/trainer.py +++ b/torchtitan/experiments/graph_trainer/trainer.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from collections.abc import Iterator from dataclasses import dataclass, field @@ -38,6 +39,7 @@ TRACE_CALL_INPUT_PREPARERS, TRACE_INPUT_PREPARERS, ) +from torchtitan.models.common.moe import uses_global_batch_lbl from torchtitan.tools.logging import logger from torchtitan.trainer import Trainer @@ -109,6 +111,14 @@ def __init__(self, config): super().__init__(config) _maybe_apply_numa_binding(self.device.index, self.device.type) + if any(uses_global_batch_lbl(m) for m in self.model_parts) and ( + self.config.compile.ep_overlap.enabled + or self.config.compile.precompile_artifact_dir + ): + # EP overlap counts tokens per chunk; a precompiled graph lacks the token-count mesh. + raise ValueError( + "load_balancing='global_aux_loss' does not support EP overlap or precompile" + ) # Lazy state for aot_fx_trace mode self._traced_step: TracedResult | None = None diff --git a/torchtitan/experiments/rl/actors/trainer.py b/torchtitan/experiments/rl/actors/trainer.py index a8aef9d7504..0673fa28087 100644 --- a/torchtitan/experiments/rl/actors/trainer.py +++ b/torchtitan/experiments/rl/actors/trainer.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import logging import os @@ -111,6 +112,7 @@ def __init__( self.config = config self.compile_config = compile_config + self.compiled_grad_clip = dist_utils.build_compiled_grad_clip(compile_config) self.loss_fn = config.loss.build() # TODO: add support to compile the loss. @@ -450,6 +452,7 @@ async def optim_step(self) -> OptimStepOutput: foreach=True, pp_mesh=self.parallel_dims.get_optional_mesh("pp"), ep_enabled=self.parallel_dims.ep_enabled, + compiled_grad_clip=self.compiled_grad_clip, ) with sl.log_trace_span("optim"): diff --git a/torchtitan/experiments/torchft/trainer.py b/torchtitan/experiments/torchft/trainer.py index 535fba07cd0..0e0ddacddf8 100644 --- a/torchtitan/experiments/torchft/trainer.py +++ b/torchtitan/experiments/torchft/trainer.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import json import os @@ -53,6 +54,7 @@ def __init__(self, config: Config): self.device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}") # Device has to be set before creating TorchFT manager. device_module.set_device(self.device) + self.compiled_grad_clip = dist_utils.build_compiled_grad_clip(config.compile) # init distributed and build meshes (FT override handles ft_manager creation) self.parallel_dims = parallel_dims = self.init_distributed() @@ -419,6 +421,7 @@ def train_step( foreach=True, pp_mesh=parallel_dims.get_optional_mesh("pp"), ep_enabled=parallel_dims.ep_enabled, + compiled_grad_clip=self.compiled_grad_clip, ) self.checkpointer.maybe_wait_for_staging() self.optimizers.step() diff --git a/torchtitan/experiments/transformers_modeling_backend/model.py b/torchtitan/experiments/transformers_modeling_backend/model.py index fe354221a93..fcb0502c900 100644 --- a/torchtitan/experiments/transformers_modeling_backend/model.py +++ b/torchtitan/experiments/transformers_modeling_backend/model.py @@ -1,13 +1,14 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import copy import importlib import math -from dataclasses import dataclass, fields +from dataclasses import dataclass, field, fields import torch from torch import nn @@ -93,6 +94,22 @@ class Config(BaseModel.Config, PretrainedConfig): HuggingFace compatibility. """ + # Ported verbatim from upstream torchtitan main, which already has it: + # https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/transformers_modeling_backend/model.py#L205 + # Redeclare PretrainedConfig dataclass fields as init=False so + # ``Configurable.__init_subclass__``'s keyword-only check passes under + # transformers v5. Drop for upstream's version at the next refresh. + transformers_version: str = field(init=False, default="") + architectures: list | None = field(init=False, default=None) + output_hidden_states: bool = field(init=False, default=False) + return_dict: bool = field(init=False, default=True) + dtype: str | None = field(init=False, default=None) + chunk_size_feed_forward: int = field(init=False, default=0) + is_encoder_decoder: bool = field(init=False, default=False) + id2label: dict | None = field(init=False, default=None) + label2id: dict | None = field(init=False, default=None) + problem_type: str | None = field(init=False, default=None) + def __init__( self, titan_dense_config, diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py index f0e6bc791c4..6428bd5c111 100644 --- a/torchtitan/models/common/attention.py +++ b/torchtitan/models/common/attention.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. # Shape suffix legend # (https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd): @@ -44,6 +45,10 @@ from torchtitan.models.common.nn_modules import Linear, RMSNorm from torchtitan.models.common.rope import RoPE +from torchtitan.protocols.checkpoint_load import ( + FusedQKVLoadTarget, + is_fused_qkv_load_planning, +) from torchtitan.protocols.module import Module from torchtitan.tools.utils import round_up @@ -750,7 +755,7 @@ def __init__(self, config: Config): @spmd.local_map( out_types=({"dp": spmd.S(0), "cp": spmd.S(1), "tp": spmd.S(2)},) * 3 ) - def forward( # pyrefly: ignore[bad-override] + def forward( self, x: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bs, seqlen, _ = x.shape @@ -808,6 +813,27 @@ def _split_qkv_on_save(module, state_dict, prefix, local_metadata) -> None: if key not in state_dict: continue tensor = state_dict.pop(key) + n_kv = tensor.shape[0] // (r * hd) + tail = (tensor.shape[1],) if ndim == 4 else () + + if is_fused_qkv_load_planning(): + projection_shapes = { + "q": torch.Size((n_kv * hpk * hd, *tail)), + "k": torch.Size((n_kv * hd, *tail)), + "v": torch.Size((n_kv * hd, *tail)), + } + for projection, shape in projection_shapes.items(): + state_dict[f"{prefix}w{projection}.{param}"] = FusedQKVLoadTarget( + fused_tensor=tensor, + projection=projection, + shape=shape, + dtype=tensor.dtype, + n_kv_heads=n_kv, + heads_per_kv=hpk, + head_dim=hd, + ) + continue + # Gather to Replicate so the n_kv-leading reshape is local (dim 0 # unsharded) when a Shard(0) split would not divide n_kv_heads # (e.g. dp_shard=8, n_kv_heads=4); stays a DTensor for the copy. @@ -815,8 +841,6 @@ def _split_qkv_on_save(module, state_dict, prefix, local_metadata) -> None: tensor = tensor.redistribute( tensor.device_mesh, [Replicate()] * tensor.device_mesh.ndim ) - n_kv = tensor.shape[0] // (r * hd) - tail = (tensor.shape[1],) if ndim == 4 else () w = tensor.reshape(n_kv, r, hd, *tail) state_dict[f"{prefix}wq.{param}"] = ( w[:, :hpk].reshape(-1, *tail).contiguous() diff --git a/torchtitan/models/common/config_utils.py b/torchtitan/models/common/config_utils.py index 6e43b08e6cd..a202229165a 100644 --- a/torchtitan/models/common/config_utils.py +++ b/torchtitan/models/common/config_utils.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. """Shared config builder helpers for model registries. @@ -23,6 +24,7 @@ FusedQKVLinear, GQAttention, QKVLinear, + ScaledDotProductAttention, VarlenAttention, ) from torchtitan.models.common.decoder import Decoder @@ -36,9 +38,11 @@ HybridEPTokenDispatcher, LocalTokenDispatcher, MinimalAsyncEPTokenDispatcher, + NeuronEPTokenDispatcher, ) from torchtitan.protocols.model_spec import ModelSpec from torchtitan.protocols.module import Module +from torchtitan.tools.logging import logger def decoder_vocab_size(model_spec: ModelSpec) -> int: @@ -53,12 +57,20 @@ def get_attention_config( ) -> Module.Config: """Map backend string to an inner_attention config. - Language models always use block_causal masking (the dataloaders always - emit per-document positions), so every backend here is a masked attention - backend. ``ScaledDotProductAttention`` only supports a boolean ``is_causal`` - flag and cannot consume per-document positions, so it is not a valid - language-model backend (it remains available for Flux, which builds it - directly). + Supported backends for language models: + + - "flex": FlexAttention with a compiled block mask. Supports causal, + packed-document, and sliding-window masking. + - "flex_flash": FlexAttention routed to the FLASH kernel backend + (Hopper/Blackwell GPUs only). + - "varlen": FA3 varlen attention. True per-document separation via + cumulative sequence metadata; supports sliding windows. + - "sdpa": F.scaled_dot_product_attention. Only exposes a boolean + ``is_causal``, so packed documents in one sequence attend across + document boundaries (a warning is logged when this backend is built). + + The gpt_oss-only "nki" backend is constructed by the gpt_oss registry + directly, not here. """ if backend == "flex": return FlexAttention.Config() @@ -75,10 +87,15 @@ def get_attention_config( elif backend == "varlen": return VarlenAttention.Config() elif backend == "sdpa": - raise ValueError( - "sdpa is no longer supported for language models; positions are " - "always available so use flex, flex_flash, or varlen." + # The text dataloaders always pack documents, and SDPA cannot express + # per-document masking, so this limitation applies to every SDPA run. + logger.warning( + "attn_backend='sdpa' cannot express per-document masking (SDPA " + "only exposes is_causal). Documents packed into the same sequence " + "will attend across document boundaries. Use 'varlen' or 'flex' " + "for per-document separation." ) + return ScaledDotProductAttention.Config() else: raise ValueError(f"Unknown backend: {backend}") @@ -320,7 +337,11 @@ def make_token_dispatcher_config( - "deepep": Uses DeepEP custom kernels for H100/NVLink Switch - "hybridep": Uses HybridEP with TMA optimization for GB200/NVLink72 - "minimal_async_ep": Uses MinimalAsyncEP for constrained DP>=EP - + - "neuron_ep": Uses the Neuron NKI all-to-all-v kernels (permute_a2av / + unpermute_a2av) for EP on Trainium; falls back to local dispatch when EP=1. + The published packed kernels send only real per-destination rows and require + Trn3 LNC=2, EP >= 8, and num_local_experts == 1 (EP == num_experts). + Multi-local-expert regrouping is not supported. DeepEP/HybridEP requires installation: https://github.com/deepseek-ai/DeepEP @@ -360,10 +381,18 @@ def make_token_dispatcher_config( num_experts=num_experts, top_k=top_k, ) + elif comm_backend == "neuron_ep": + # Neuron NKI all-to-all-v dispatcher (permute_a2av / unpermute_a2av): + # packed token-major transmission for num_local_experts == 1. + return NeuronEPTokenDispatcher.Config( + num_experts=num_experts, + top_k=top_k, + ) else: raise ValueError( f"Unknown comm_backend: '{comm_backend}'. " - "Must be one of 'standard', 'deepep', 'hybridep', 'minimal_async_ep'." + "Must be one of 'standard', 'deepep', 'hybridep', 'minimal_async_ep', " + "'neuron_ep'." ) diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index e170150410a..b5592fa8b17 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import dataclasses from dataclasses import dataclass @@ -23,6 +24,7 @@ FlexAttention, get_causal_mask_mod, get_efficient_causal_mask_mod_for_packed_document, + ScaledDotProductAttention, VarlenAttention, ) from torchtitan.models.common.embedding import Embedding @@ -325,6 +327,10 @@ def get_attention_masks( return self._create_flex_attention_mask_for_document(positions, attn_config) elif isinstance(inner_attn, VarlenAttention.Config): return create_varlen_metadata_for_document(positions) + elif isinstance(inner_attn, ScaledDotProductAttention.Config): + # SDPA uses is_causal only; no explicit masks. Packed documents + # attend across boundaries (get_attention_config warns). + return None else: raise TypeError( f"Only VarlenAttention and FlexAttention support attention masks, " diff --git a/torchtitan/models/common/moe.py b/torchtitan/models/common/moe.py index 4c3167510eb..6b7389cd33b 100644 --- a/torchtitan/models/common/moe.py +++ b/torchtitan/models/common/moe.py @@ -1,17 +1,19 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from dataclasses import dataclass -from typing import Literal +from typing import cast, Literal, Protocol import spmd_types as spmd - import torch +import torch.distributed._functional_collectives as funcol import torch.nn.functional as F from torch import nn +from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor from torchtitan.distributed.spmd_types import maybe_set_sparse_mesh, spmd_mesh_size @@ -22,6 +24,7 @@ from .token_dispatcher import DeepEPTokenDispatcher, LocalTokenDispatcher + # Shape suffix legend # (https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd): # B = batch, L = sequence length, D = model dimension, @@ -350,6 +353,9 @@ class Config(Module.Config): experts: GroupedExperts.Config router: TokenChoiceTopKRouter.Config load_balance_coeff: float | None = 1e-3 + # "expert_bias": aux-loss-free bias (default). "global_aux_loss": Qwen3's global-batch + # load-balancing loss with coefficient load_balance_coeff (see attach_global_batch_lbl). + load_balancing: Literal["expert_bias", "global_aux_loss"] = "expert_bias" shared_experts: FeedForward.Config | None = None # TODO(pianpwk): Remove this once MoE combine can derive the local # sequence shape directly from the input layout. @@ -371,8 +377,15 @@ def __init__(self, config: Config): # expert_bias_E is updated outside the model in an optimizer step pre hook # to work with gradient accumulation. self.load_balance_coeff = config.load_balance_coeff + self.load_balancing = config.load_balancing + if self.load_balancing not in ("expert_bias", "global_aux_loss"): + raise ValueError(f"unknown load_balancing {self.load_balancing!r}") + # Set by the trainer for the global-batch loss (see attach_global_batch_lbl). + self.lbl_batch_mesh: DeviceMesh | None = None + self.lbl_num_microbatches: int = 1 if self.load_balance_coeff is not None: assert self.load_balance_coeff > 0.0 + if self.load_balance_coeff is not None and self.load_balancing == "expert_bias": self.register_buffer( "expert_bias_E", torch.zeros(num_experts, dtype=torch.float32), @@ -463,6 +476,11 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: # effect on the expert bias update thanks to the torch.sign() operator. with torch.no_grad(): self.tokens_per_expert_E.add_(num_local_tokens_per_expert_E) + lbl = ( + self.load_balancing == "global_aux_loss" + and self.load_balance_coeff is not None + and torch.is_grad_enabled() + ) out_BLD = self.experts( x_BLD, @@ -507,6 +525,14 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: # Drop the tokens padded on for SP sharding, restoring (B, L, D). out_BLD = out_BLD[:, : L - seq_pad, :] # --------------------------------------------------------------------- + if lbl: + out_BLD = attach_global_batch_lbl( + out_BLD, + num_local_tokens_per_expert_E, + scores_BLE.sum(dim=(0, 1)), + self, + top_k=self.router.top_k, + ) return out_BLD def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: @@ -519,7 +545,91 @@ def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> No self.tokens_per_expert_E = torch.zeros( self.experts.num_experts, dtype=torch.float32 ) - if self.load_balance_coeff is not None: + if self.expert_bias_E is not None: self.expert_bias_E = torch.zeros( self.experts.num_experts, dtype=torch.float32 ) + + +class _AttachAuxLoss(torch.autograd.Function): + """Identity on ``x``; backward seeds ``aux`` with 1 inside the caller's own backward.""" + + @staticmethod + # pyrefly: ignore [bad-override] + def forward(ctx, x, aux): + ctx.save_for_backward(aux) + return x + + @staticmethod + # pyrefly: ignore [bad-override] + def backward(ctx, grad_x): + (aux,) = ctx.saved_tensors + return grad_x, torch.ones_like(aux) + + +def uses_global_batch_lbl(model: nn.Module) -> bool: + return any( + getattr(m, "load_balancing", None) == "global_aux_loss" + and m.load_balance_coeff is not None + for m in model.modules() + ) + + +class _GlobalBatchLBLModule(Protocol): + """The state ``attach_global_batch_lbl`` reads from an MoE layer (MoE or a backend MoE).""" + + load_balance_coeff: float | None + lbl_batch_mesh: DeviceMesh | None + lbl_num_microbatches: int + + +def configure_global_batch_lbl( + model_parts: list[nn.Module], + *, + batch_mesh: DeviceMesh | None, + num_forwards: int, +) -> None: + """Give each global-batch-LBL MoE its token-count group and the step's forward count.""" + for part in model_parts: + for m in part.modules(): + if getattr(m, "load_balancing", None) == "global_aux_loss": + lbl_module = cast(_GlobalBatchLBLModule, m) + lbl_module.lbl_batch_mesh = batch_mesh + lbl_module.lbl_num_microbatches = num_forwards + + +def attach_global_batch_lbl( + x: torch.Tensor, + counts_E: torch.Tensor, + prob_mass_E: torch.Tensor, + moe: _GlobalBatchLBLModule, + *, + top_k: int, +) -> torch.Tensor: + """Attach Qwen3's global-batch load-balancing loss (arXiv 2501.11873) to ``x``. + + ``coeff * E * sum_i f_i * P_i``: ``f_i`` from token counts summed over every rank holding + distinct tokens (``moe.lbl_batch_mesh``), ``P_i`` this rank's probability mass over the step's + global token count, so ranks sum to the global-batch loss and micro-batches share the step's + normalization like the main loss. + The token count is every routed token, as in Megatron's aux loss (with masked labels it exceeds + the main loss's valid-token count). It is attached rather than added to the main loss so its gradient flows inside this layer's + backward: a loss-level term reaches every layer at once and, under activation checkpointing, + recomputes every block at once (2x peak memory). Each micro-batch uses its own global counts + (Megatron's running mean is a follow-up). + """ + counts_E, prob_mass_E = ( + t.full_tensor() if isinstance(t, DTensor) else t + for t in (counts_E.detach().float(), prob_mass_E) + ) + if moe.lbl_batch_mesh is not None: + counts_E = funcol.all_reduce(counts_E, "sum", moe.lbl_batch_mesh) + global_tokens = counts_E.sum() / top_k * moe.lbl_num_microbatches + assert moe.load_balance_coeff is not None + aux = ( + moe.load_balance_coeff + * counts_E.numel() + * (counts_E / counts_E.sum() * prob_mass_E).sum() + / global_tokens + ) + return _AttachAuxLoss.apply(x, aux) diff --git a/torchtitan/models/common/moe_sharding.py b/torchtitan/models/common/moe_sharding.py index a89f6308d61..7974112268f 100644 --- a/torchtitan/models/common/moe_sharding.py +++ b/torchtitan/models/common/moe_sharding.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. """Config-based sharding helpers for MoE submodules.""" @@ -318,8 +319,6 @@ def set_moe_sharding_config( # num_local_tokens_per_expert_E is routing metadata, but it is # still a DTensor input to local_map and must have placements. _tokens_per_expert_placement(enable_ep=enable_ep), - ) - if enable_ep - else None, + ), ), ) diff --git a/torchtitan/models/common/rope.py b/torchtitan/models/common/rope.py index b1282915229..914c7b929f3 100644 --- a/torchtitan/models/common/rope.py +++ b/torchtitan/models/common/rope.py @@ -1,8 +1,12 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. + +# Shape suffixes: B = batch, L = sequence, N = heads, H = head dimension, +# C = cache width (H / 2 for complex caches, 2 * H for cos/sin caches). import math from dataclasses import dataclass @@ -47,6 +51,10 @@ class RoPE(Module): cosine/sine caches. """ + _cp_rank: int = 0 + _cp_size: int = 1 + _cp_load_balancer: str | None = None + @dataclass(kw_only=True, slots=True) class Config(Module.Config): dim: int @@ -64,12 +72,19 @@ class Config(Module.Config): beta_slow: float = 1.0 original_seq_len: int = 4096 mscale: float = 0.0 + use_prefix_cache_for_positions: bool = False + # For causal, non-packed training with sequential position IDs, build + # the cache from the sequence prefix instead of gathering cache rows. def __init__(self, config: Config): super().__init__() self.config = config self.register_buffer("cache", self._precompute_cache(), persistent=False) + @property + def uses_prefix_cache_for_positions(self) -> bool: + return self.config.use_prefix_cache_for_positions + def _precompute_cache(self) -> torch.Tensor: """Build the reusable cache for all positions up to ``max_seq_len``. @@ -78,6 +93,41 @@ def _precompute_cache(self) -> torch.Tensor: """ raise NotImplementedError + def _reshape_prefix_cache(self, query_BLNH: torch.Tensor) -> torch.Tensor: + """Select sequential cache rows for this rank's context-parallel shard.""" + if self._cp_size == 1: + return _reshape_for_broadcast(self.cache, query_BLNH.shape, None) + if isinstance(query_BLNH, DTensor): + if not isinstance(self.cache, DTensor) or not all( + isinstance(p, Replicate) for p in self.cache.placements + ): + raise ValueError( + "DTensor queries require a replicated DTensor RoPE cache" + ) + # The cache is broadcast over batch and heads, and only its + # sequence dimension follows the query's sharding. + placements = tuple( + p if isinstance(p, Shard) and p.dim == 1 else Replicate() + for p in query_BLNH.placements + ) + cache_1L1C = _reshape_prefix_cache( + self.cache.to_local(), + query_BLNH.to_local(), + self._cp_rank, + self._cp_size, + self._cp_load_balancer, + ) + return DTensor.from_local( + cache_1L1C, query_BLNH.device_mesh, placements, run_check=False + ) + return _reshape_prefix_cache( + self.cache, + query_BLNH, + self._cp_rank, + self._cp_size, + self._cp_load_balancer, + ) + def _reshape_cache( self, query: torch.Tensor, @@ -134,10 +184,15 @@ def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> No self.cache = self._precompute_cache() +@dataclass(kw_only=True, slots=True) +class ComplexRoPEConfig(RoPE.Config): + pass + + class ComplexRoPE(RoPE): - @dataclass(kw_only=True, slots=True) - class Config(RoPE.Config): - pass + # Keep the public ``ComplexRoPE.Config`` API, while giving torch.compile's + # caching precompile guard serializer a top-level dataclass type. + Config = ComplexRoPEConfig def _precompute_cache(self) -> torch.Tensor: """Precompute complex cis values. @@ -232,12 +287,14 @@ def _reshape_cache( Returns: Cache of shape ``(1 or batch, seq_len, 1, dim / 2)``. """ - positions = _maybe_wrap_positions(positions, query) - if positions is not None: - _maybe_check_max_pos(positions, max_valid_pos=self.cache.shape[0] - 1) # Complex RoPE cache has width dim / 2 because each complex value # represents a pair of real dimensions. complex_query_shape = (*query.shape[:-1], query.shape[-1] // 2) + if self.uses_prefix_cache_for_positions: + return self._reshape_prefix_cache(query) + positions = _maybe_wrap_positions(positions, query) + if positions is not None: + _maybe_check_max_pos(positions, max_valid_pos=self.cache.shape[0] - 1) return _reshape_for_broadcast(self.cache, complex_query_shape, positions) @staticmethod @@ -254,10 +311,15 @@ def apply_rotary_emb( return xq_out.type_as(query), xk_out.type_as(key) +@dataclass(kw_only=True, slots=True) +class CosSinRoPEConfig(RoPE.Config): + pass + + class CosSinRoPE(RoPE): - @dataclass(kw_only=True, slots=True) - class Config(RoPE.Config): - pass + # Keep the public ``CosSinRoPE.Config`` API, while giving torch.compile's + # caching precompile guard serializer a top-level dataclass type. + Config = CosSinRoPEConfig def _precompute_cache(self) -> torch.Tensor: """Precompute cos/sin values. @@ -324,6 +386,8 @@ def _reshape_cache( Returns: Cache of shape ``(1 or batch, seq_len, 1, dim * 2)``. """ + if self.uses_prefix_cache_for_positions: + return self._reshape_prefix_cache(query) positions = _maybe_wrap_positions(positions, query) if positions is not None: _maybe_check_max_pos(positions, max_valid_pos=self.cache.shape[0] - 1) @@ -352,6 +416,40 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x2, x1), dim=-1) +@spmd.local_map(out_types={"dp": spmd.R, "cp": spmd.S(1), "tp": spmd.R}) +def _reshape_prefix_cache( + cache_LC: torch.Tensor, + query_BLNH: torch.Tensor, + cp_rank: int, + cp_size: int, + load_balancer: str | None, +) -> torch.Tensor: + """Build a broadcast cache without gathering device-side position IDs.""" + seq_len = query_BLNH.shape[1] + if seq_len * cp_size > cache_LC.shape[0]: + raise ValueError("The context-parallel sequence exceeds the RoPE cache") + if load_balancer is None: + start = cp_rank * seq_len + cache_LC = cache_LC[start : start + seq_len] + elif load_balancer == "headtail": + if seq_len % 2: + raise ValueError( + "Head-tail context parallelism requires even shard lengths" + ) + half = seq_len // 2 + head = cp_rank * half + tail = (2 * cp_size - cp_rank - 1) * half + cache_LC = torch.cat( + (cache_LC[head : head + half], cache_LC[tail : tail + half]) + ) + else: + raise ValueError( + "Prefix RoPE caches require context_parallel_load_balancer=None " + "or 'headtail'" + ) + return cache_LC.view(1, seq_len, 1, cache_LC.shape[-1]) + + @spmd.local_map(out_types={"dp": spmd.S(0), "cp": spmd.S(1), "tp": spmd.R}) def _reshape_for_broadcast( rope_cache: torch.Tensor, diff --git a/torchtitan/models/common/token_dispatcher.py b/torchtitan/models/common/token_dispatcher.py index 859e2eefafe..f68429c78ac 100644 --- a/torchtitan/models/common/token_dispatcher.py +++ b/torchtitan/models/common/token_dispatcher.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from dataclasses import dataclass from typing import cast, ClassVar @@ -19,6 +20,13 @@ init_buffer as minimal_async_ep_init_buffer, MinimalAsyncEPDispatchMetadata, ) +from torchtitan.distributed.neuron_ep import ( + build_ep_replica_group_ranks, + build_packed_source_indices, + gmm_align, + permute_a2av_dispatch_compact, + unpermute_a2av_combine_packed, +) from torchtitan.distributed.spmd_types import current_spmd_mesh, maybe_set_sparse_mesh from torchtitan.distributed.utils import get_spmd_backend from torchtitan.ops.scatter_add import deterministic_scatter_add @@ -43,6 +51,26 @@ class AllToAllDispatchMetadata(LocalDispatchMetadata): output_splits: list[int] +@dataclass(frozen=True, kw_only=True) +class NeuronEPDispatchMetadata(LocalDispatchMetadata): + """Metadata returned by NeuronEPTokenDispatcher.dispatch() for use in combine(). + + ``token_indices_experts_sorted_N`` / ``topk_scores_experts_sorted_N`` are the + pre-expanded (N = T*top_k) per-row token-id and gate-score maps, index-aligned + with the kernel's ``routed_input_ND`` rows so combine can apply scores and + scatter-add back to T tokens. They are TOKEN-MAJOR (``arange(N)//top_k`` and + ``topk_scores_TK.reshape(-1)``): the dispatcher leaves the expert-sort to the + kernel's packed-index permute, so combine relies only on 1:1 alignment with + the kernel row order, not on the ordering being expert-sorted. (The field names + keep the ``_experts_sorted_N`` suffix for historical/dataclass compatibility.) + """ + + send_counts_1E: torch.Tensor # noqa: N815 # int32 [1, EP], rows sent per dest + recv_counts_1E: torch.Tensor # noqa: N815 # int32 [1, EP], rows received per source + compact_capacity: int # row count of the packed [C,H] expert buffer + packed_source_indices_N: torch.Tensor # noqa: N815 + + class LocalTokenDispatcher(Configurable): """Token dispatcher for EP=1. Handles local token reordering only. @@ -746,6 +774,321 @@ def _unpermute(self, routed_output_RD, input_shape, permuted_indices): return out_unpermuted_RD[:-1] +class NeuronEPTokenDispatcher(BaseEPTokenDispatcher): + """EP token dispatcher backed by Neuron NKI all-to-all-v kernels. + + Design: + + 1. Packed capacity. The supported EP path requires + ``num_local_experts == 1`` and sends real per-destination counts into a + packed ``[C,H]`` receive buffer. The aligned capacity + ``C = FSDP * local_batch * global_sequence_length`` is the unique-token + count represented by the EP group. With multiple local experts, capacity + scales by ``min(top_k, num_local_experts)`` once that path is enabled. + The kernel output is ``[C,H]``, so its element count is ``C * H``. Its + NKI kernel overlaps the count exchange with local payload packing. + + 2. Token-major kernel input. Feeds the NKI ``permute_a2av`` kernel a TOKEN-MAJOR + ``routed_input_ND`` (row ``r = t*K + k`` = token ``t`` replicated for its k-th + expert choice) and lets the kernel's packed vector do the routing sort + in a single pass -- no glue-side argsort + gather. + + The published kernels currently require Trn3, LNC=2, EP >= 8, and complete + four-rank device groups. Multi-local-expert packed regrouping is a separate + milestone because source-major receive counts do not identify local experts. + + Scope: EP with ``num_experts % ep_size == 0``. TP/SP (``sp_size > 1``) is supported + -- routing is unaffected by the sequence shard (the packed vector keys off expert ids + only, and the a2a-v round-trip preserves per-row identity), so SP only offsets the + final scatter in ``combine`` to global token positions. CP is not supported. + + Tensor shape suffixes (module-local): ``N`` routed rows (T*top_k), ``D``/``H`` hidden, + ``E`` EP degree, ``e`` local experts per rank, and ``C`` packed receive capacity. + """ + + @dataclass(kw_only=True, slots=True) + class Config(BaseEPTokenDispatcher.Config): + # Optional per-source capacity override. The packed buffer has + # EP*min(top_k, num_local_experts)*capacity_per_expert rows when set. + capacity_per_expert: int | None = None + + def __init__(self, config: Config): + super().__init__(config) + self._cap_override = config.capacity_per_expert + # Full SPMD replica-group rank lists for nki.collectives.ReplicaGroup, + # set in wire_meshes once the EP mesh is known. + self._ep_ranks: list[list[int]] | None = None + + def wire_meshes(self, *, ep_mesh, tp_mesh) -> None: + super().wire_meshes(ep_mesh=ep_mesh, tp_mesh=tp_mesh) + if ep_mesh is not None: + if self.num_experts % ep_mesh.size() != 0: + raise ValueError( + f"neuron_ep requires num_experts ({self.num_experts}) divisible by " + f"ep_size ({ep_mesh.size()})." + ) + self._ep_ranks = build_ep_replica_group_ranks(ep_mesh) + + def _packed_capacity( + self, + local_token_count: int, + ep_size: int, + num_local_experts: int, + ) -> int: + """Return the grouped-mm-aligned dropless receive capacity. + + ``local_token_count * ep_size`` equals + ``FSDP * local_batch * global_sequence_length`` for the EP/TP mesh. + A token can target at most ``min(top_k, num_local_experts)`` experts on + one destination rank because its top-k expert IDs are distinct. + """ + alignment = gmm_align() + expert_multiplier = min(self.top_k, num_local_experts) + model_capacity = local_token_count * ep_size * expert_multiplier + capacity = model_capacity + if self._cap_override is not None: + capacity = ep_size * expert_multiplier * self._cap_override + if capacity < model_capacity: + raise ValueError( + "compact neuron_ep requires " + "EP*min(top_k,num_local_experts)*capacity_per_expert " + ">= model capacity " + f"({ep_size}*{expert_multiplier}*{self._cap_override}=" + f"{capacity} < {model_capacity})." + ) + return ((capacity + alignment - 1) // alignment) * alignment + + def _build_packed_routing( + self, + topk_expert_ids_TK: torch.Tensor, + num_local_tokens_per_expert_E: torch.Tensor, + ): + """Build stable destination-major routing for one local expert per rank. + + Returns ``(counts[1, EP], source_indices[N])``. Peer segments preserve + token-major arrival order, with every routed row appearing exactly once. + """ + # The router already counts local tokens for each expert. With one expert + # per rank and distinct top-k choices, these are exactly the peer counts. + send_counts_1E = num_local_tokens_per_expert_E.to(torch.int32).view(1, -1) + packed_source_indices_N = build_packed_source_indices( + topk_expert_ids_TK.reshape(-1), send_counts_1E + ) + return send_counts_1E, packed_source_indices_N + + def _dispatch_compact( + self, + routed_input_ND: torch.Tensor, + topk_expert_ids_TK: torch.Tensor, + token_indices_per_row_N: torch.Tensor, + topk_scores_per_row_N: torch.Tensor, + num_local_tokens_per_expert_E: torch.Tensor, + recv_capacity: int, + ) -> tuple[torch.Tensor, torch.Tensor, NeuronEPDispatchMetadata]: + """num_local_experts == 1 compact dispatch: send only real rows. + + Builds a contiguous destination-major source-index vector. The packed kernel issues the + count exchange before local packing and delays synchronization until the + received counts are consumed, permitting the compiler to overlap them. It + receives directly into the static ``[recv_capacity, H]`` expert buffer. + """ + device = routed_input_ND.device + send_counts_1E, packed_source_indices_N = self._build_packed_routing( + topk_expert_ids_TK, num_local_tokens_per_expert_E + ) + # recv_counts is a kernel OUTPUT (variable-count a2a-v derives it internally). + routed_input_RD, recv_counts_1E = permute_a2av_dispatch_compact( + routed_input_ND, + packed_source_indices_N, + send_counts_1E, + recv_capacity, + self._ep_ranks, + ) + padded_counts_e = torch.full( + (1,), recv_capacity, dtype=torch.int32, device=device + ) + metadata = NeuronEPDispatchMetadata( + token_indices_experts_sorted_N=token_indices_per_row_N, + topk_scores_experts_sorted_N=topk_scores_per_row_N, + send_counts_1E=send_counts_1E, + recv_counts_1E=recv_counts_1E, + compact_capacity=recv_capacity, + packed_source_indices_N=packed_source_indices_N, + ) + return routed_input_RD, padded_counts_e, metadata + + # pyrefly: ignore [bad-override] + def dispatch( + self, + x_TD: torch.Tensor, + topk_scores_TK: torch.Tensor, + topk_expert_ids_TK: torch.Tensor, + num_local_tokens_per_expert_E: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, NeuronEPDispatchMetadata]: + """Packed dispatch through the published NKILib A2AV API. + + The glue does not expert-sort on the send side. We expand ``x_TD`` to + token-major ``routed_input_ND`` (row r = t*K + k = token t replicated for + its k-th expert) and let the kernel gather rows from the packed vector. + """ + if self.ep_mesh is None: + # Local (EP=1) fallback: unchanged. LocalTokenDispatcher.dispatch still + # does expert-sort on the send side and consumes num_local_tokens_per_expert_E + # verbatim; only the EP path benefits from the token-major kernel-permute layout. + return LocalTokenDispatcher.dispatch( + self, + x_TD, + topk_scores_TK, + topk_expert_ids_TK, + num_local_tokens_per_expert_E, + ) + assert self._ep_ranks is not None + ep_size = self.ep_mesh.size() + num_local_experts = self.num_experts // ep_size + t = x_TD.shape[0] + k = topk_expert_ids_TK.shape[1] + if topk_expert_ids_TK.shape != topk_scores_TK.shape: + raise ValueError( + "neuron_ep requires matching top-k score and expert-id shapes; " + f"got {tuple(topk_scores_TK.shape)} and " + f"{tuple(topk_expert_ids_TK.shape)}." + ) + expected_routing_shape = (t, self.top_k) + if topk_expert_ids_TK.shape != expected_routing_shape: + raise ValueError( + "neuron_ep requires routing shape " + f"{expected_routing_shape}. Got {tuple(topk_expert_ids_TK.shape)}." + ) + if ep_size < 8 or ep_size % 4 != 0: + raise ValueError( + "packed neuron_ep requires ep_size >= 8 in complete four-rank " + f"Trn3 device groups. Got {ep_size}." + ) + n = t * k + if num_local_experts != 1: + raise ValueError( + "packed neuron_ep currently requires num_local_experts == 1; " + f"got num_experts={self.num_experts}, ep_size={ep_size}, " + f"num_local_experts={num_local_experts}." + ) + + # Token-major expansion: row r = t*K + k is token t's embedding, repeated K + # times contiguously. Only routing indices are sorted; the kernel gathers + # the hidden-state rows via packed source indices. + routed_input_ND = x_TD.repeat_interleave(k, dim=0) + # Per-row metadata index-aligned with routed_input_ND rows: + # token_indices_per_row_N[r] = t = origin token id (for combine scatter_add) + # topk_scores_per_row_N[r] = s_tk = gate score for that (token, expert) pair + # Kept in the NeuronEPDispatchMetadata fields named ``*_experts_sorted_N`` (a + # historical name): the data is token-major, not expert-sorted, but combine + # only needs 1:1 alignment with the kernel row order. + arange_N = torch.arange(n, device=x_TD.device) + token_indices_per_row_N = arange_N // k + topk_scores_per_row_N = topk_scores_TK.reshape(-1) + + # The kernel's per-source recv_counts are also per-(source, expert) when + # num_local_experts == 1. + return self._dispatch_compact( + routed_input_ND, + topk_expert_ids_TK, + token_indices_per_row_N, + topk_scores_per_row_N, + num_local_tokens_per_expert_E, + self._packed_capacity(t, ep_size, num_local_experts), + ) + + # pyrefly: ignore [bad-override] + def combine( + self, + routed_output_RD: torch.Tensor, + metadata: NeuronEPDispatchMetadata, + x_TD: torch.Tensor, + *, + num_local_tokens_after_padding: int, + local_seq_len_after_padding: int, + ) -> torch.Tensor: + """Run packed NKI combine, apply scores, and reduce routed rows N -> tokens T. + + The EP path consumes the packed ``[C,H]`` expert output directly. + Metadata's ``*_experts_sorted_N`` fields carry token-major (per-row) + values -- see ``dispatch`` -- but combine treats them as opaque + row-aligned indices/scores. + + With SP (``sp_size > 1``) each rank holds only a ``1/sp_size`` slice of the + sequence, so ``token_indices_experts_sorted_N`` are LOCAL token ids. They are + offset to GLOBAL positions here and scattered into a global-length buffer, so + every SP rank writes a DISJOINT slice (zeros elsewhere). The MoE forward output + is declared ``Partial(tp) -> Shard(seq)`` in ``set_moe_sharding_config``, and + that DTensor redistribution reduce-scatters these disjoint slices back into a + per-rank sequence shard -- which is why the buffer must be global-length and why + no TP collective belongs here. This mirrors ``AllToAllTokenDispatcher.combine`` + (the SP remap is combine-only in both: dispatch keeps indices local because it + indexes local tensors). + """ + if self.ep_mesh is None: + return LocalTokenDispatcher.combine( + self, + routed_output_RD, + metadata, + x_TD, + num_local_tokens_after_padding=num_local_tokens_after_padding, + local_seq_len_after_padding=local_seq_len_after_padding, + ) + assert self._ep_ranks is not None + + if routed_output_RD.shape[0] != metadata.compact_capacity: + raise ValueError( + "neuron_ep compact combine expects expert output with " + f"capacity ({metadata.compact_capacity}) rows. " + f"got {routed_output_RD.shape[0]}." + ) + combined_ND = unpermute_a2av_combine_packed( + routed_output_RD, + metadata.packed_source_indices_N, + metadata.send_counts_1E, + metadata.recv_counts_1E, + self._ep_ranks, + ) + + combined_ND = ( + combined_ND.to(torch.float32) + * metadata.topk_scores_experts_sorted_N.reshape(-1, 1) + ).to(x_TD.dtype) + dim = x_TD.shape[-1] + if self.sp_size > 1: + # Global-length buffer: this rank fills only its own disjoint slice (via the + # sp_rank offset below); the downstream Partial -> Shard(seq) redistribution + # reassembles all SP ranks' slices. moe.py pads the sequence to a multiple of + # sp_size before dispatch, which keeps the per-rank token count uniform -- + # the packed receive layout relies on that uniformity. + out_TD = torch.zeros( + num_local_tokens_after_padding * self.sp_size, + dim, + device=x_TD.device, + dtype=x_TD.dtype, + ) + token_indices_N = self._sp_global_token_indices( + metadata.token_indices_experts_sorted_N, + local_seq_len_after_padding, + ) + else: + out_TD = torch.zeros_like(x_TD) + token_indices_N = metadata.token_indices_experts_sorted_N + # TODO: Replace this with a fixed-order NKI token-reduction kernel + # that provides deterministic scatter-add semantics on Neuron. + # Native Neuron scatter-add is compilable but does not guarantee reduction order. + # Functionally equivalent to Tensor.scatter_add in eager mode. Use the + # explicit ATen overload so AOT capture preserves a Torch-MLIR-recognized + # operator target. + out_TD = torch.ops.aten.scatter_add.default( + out_TD, + 0, + token_indices_N.reshape(-1, 1).expand(-1, dim), + combined_ND, + ) + return out_TD + + @dataclass(frozen=True, kw_only=True) class DeepEPDispatchMetadata: """Metadata for DeepEP, HybridEP, and MinimalAsyncEP token dispatch.""" diff --git a/torchtitan/models/deepseek_v3/sharding.py b/torchtitan/models/deepseek_v3/sharding.py index 22ec5ab403a..c2093563c4b 100644 --- a/torchtitan/models/deepseek_v3/sharding.py +++ b/torchtitan/models/deepseek_v3/sharding.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from typing import TYPE_CHECKING @@ -135,6 +136,9 @@ def _set_deepseek_v3_layer_sharding( # MoE FFN (MoE-enabled layers only). if layer_cfg.moe is not None: + if hasattr(layer_cfg.moe, "set_sharding_config"): + layer_cfg.moe.set_sharding_config(enable_sp=enable_sp, enable_ep=enable_ep) + return set_moe_sharding_config( layer_cfg.moe, enable_ep=enable_ep, diff --git a/torchtitan/models/flux/trainer.py b/torchtitan/models/flux/trainer.py index 886f41831a9..a055d14d295 100644 --- a/torchtitan/models/flux/trainer.py +++ b/torchtitan/models/flux/trainer.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import time from collections.abc import Iterable, Iterator @@ -316,6 +317,7 @@ def train_step( foreach=True, pp_mesh=parallel_dims.get_optional_mesh("pp"), ep_enabled=parallel_dims.ep_enabled, + compiled_grad_clip=self.compiled_grad_clip, ) self.checkpointer.maybe_wait_for_staging() self.optimizers.step() diff --git a/torchtitan/models/gpt_oss/__init__.py b/torchtitan/models/gpt_oss/__init__.py index 37e1b98bc5a..564494fce37 100644 --- a/torchtitan/models/gpt_oss/__init__.py +++ b/torchtitan/models/gpt_oss/__init__.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import dataclasses from collections.abc import Callable @@ -23,6 +24,7 @@ from torchtitan.models.common.attention import ( FusedQKVLinear, QKVLinear, + ScaledDotProductAttention, VarlenAttention, ) from torchtitan.models.common.config_utils import ( @@ -89,7 +91,17 @@ def _make_gptoss_attn_config( Sinks also use depth-scaled init. """ - inner_attention = get_attention_config(attn_backend) + if attn_backend == "nki": + from torchtitan_neuron.models.gpt_oss.nki_attention import ( + NKIFlashAttentionWrapper, + ) + + inner_attention = NKIFlashAttentionWrapper.Config( + softmax_scale=head_dim**-0.5, + sliding_window=sliding_window_size, + ) + else: + inner_attention = get_attention_config(attn_backend) if sliding_window_size is not None and isinstance( inner_attention, VarlenAttention.Config @@ -98,6 +110,17 @@ def _make_gptoss_attn_config( inner_attention, window_size=(sliding_window_size - 1, 0) ) + if sliding_window_size is not None and isinstance( + inner_attention, ScaledDotProductAttention.Config + ): + # Only varlen (window_size) and flex (mask mod) can express the + # window; SDPA would silently run full attention on window layers. + raise ValueError( + "attn_backend='sdpa' cannot express sliding-window attention " + f"(sliding_window_size={sliding_window_size}); use 'varlen', " + "'flex', or 'nki'" + ) + sinks_init = { "sinks": partial(nn.init.trunc_normal_, std=depth_scaled_std(0.02, layer_id)) } diff --git a/torchtitan/models/gpt_oss/model.py b/torchtitan/models/gpt_oss/model.py index fd2a6d46d8f..f750f935e51 100644 --- a/torchtitan/models/gpt_oss/model.py +++ b/torchtitan/models/gpt_oss/model.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import dataclasses import math @@ -33,6 +34,18 @@ from torchtitan.protocols.module import Module +class _ForwardOnlyScale(torch.autograd.Function): + """y = x * scale in forward; gradient passes through unscaled.""" + + @staticmethod + def forward(ctx, x, scale): + return x * scale + + @staticmethod + def backward(ctx, grad): + return grad, None + + class ScaledBiasRowwiseLinear(Linear): """ Rowwise linear whose local bias contribution is scaled by TP degree. @@ -60,7 +73,11 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight ) bias = self.bias.to_local() if isinstance(self.bias, DTensor) else self.bias - bias = bias / self.tp_degree + # Each TP rank adds bias/TP to its Partial output so the TP sum carries the + # bias exactly once. The scale must be forward-only: the bias is stored + # Replicate on TP and its gradient is never summed across ranks, so each + # rank must hold the full (identical) bias gradient, not grad/TP. + bias = _ForwardOnlyScale.apply(bias, 1.0 / self.tp_degree) return F.linear(input, weight, bias) @@ -134,15 +151,20 @@ def forward( q, k = self.rope(q, k, positions) - output = self.inner_attention( - q, - k, - v, - attention_masks=attention_masks, - scale=self.softmax_scale, - enable_gqa=self.enable_gqa, - out_transform=self._apply_sinks, - ) + # Upstream generic attention applies GPT-OSS sinks through out_transform. + # NKI consumes sinks inside the fused kernel so backward can emit grad_sinks. + if getattr(self.inner_attention, "uses_attention_sinks_natively", False): + output = self.inner_attention(q, k, v, self.sinks) + else: + output = self.inner_attention( + q, + k, + v, + attention_masks=attention_masks, + scale=self.softmax_scale, + enable_gqa=self.enable_gqa, + out_transform=self._apply_sinks, + ) # Reshape and project output output = output.reshape( @@ -230,7 +252,11 @@ def update_from_config( config, **kwargs, ) -> None: - Decoder.Config.update_from_config(self, config=config, **kwargs) + if self._has_dispatcherless_moe(): + self._update_dispatcherless_moe_from_config(config=config, **kwargs) + else: + Decoder.Config.update_from_config(self, config=config, **kwargs) + parallelism = config.parallelism from torchtitan.models.gpt_oss.sharding import set_gpt_oss_sharding_config @@ -241,6 +267,87 @@ def update_from_config( enable_ep=parallelism.expert_parallel_degree > 1, ) + def _has_dispatcherless_moe(self) -> bool: + return any( + layer.moe is not None + and not hasattr(layer.moe.experts, "token_dispatcher") + for layer in self.layers + ) + + def _update_dispatcherless_moe_from_config(self, *, config, **kwargs) -> None: + """Apply common decoder runtime updates for GPT-OSS NKI MoE. + + GPT-OSS NKI experts run without a token dispatcher. Keep that + exception local to this model instead of weakening the common decoder + and minimal-async EP validation paths for all MoE models. + """ + from torchtitan.config import ParallelismConfig + from torchtitan.trainer import Trainer + + assert hasattr(config, "parallelism"), ( + "config passed to update_from_config must provide " + "a parallelism field." + ) + parallelism = config.parallelism + assert isinstance(parallelism, ParallelismConfig), ( + "config.parallelism must be a ParallelismConfig, got " + f"{type(parallelism).__name__}." + ) + + if self.enable_weight_tying and parallelism.pipeline_parallel_degree > 1: + raise NotImplementedError( + "Weight tying is not supported with Pipeline Parallel." + ) + + if parallelism.pipeline_parallel_degree > 1 and any( + layer.attention is not None + and isinstance(layer.attention.inner_attention, VarlenAttention.Config) + for layer in self.layers + ): + raise ValueError( + "Pipeline Parallel is not compatible with VarlenAttention. " + "Use a FlexAttention backend (attn_backend='flex' or " + "'flex_flash') for pipelined models." + ) + + tp = parallelism.tensor_parallel_degree + attention = self.first_attention + if tp > 1 and attention is not None: + n_heads = attention.n_heads + n_kv_heads = getattr(attention, "n_kv_heads", None) or n_heads + if n_heads % tp != 0: + raise ValueError( + f"tensor_parallel_degree ({tp}) must divide " + f"n_heads ({n_heads})." + ) + if n_kv_heads % tp != 0: + raise ValueError( + f"tensor_parallel_degree ({tp}) must divide " + f"n_kv_heads ({n_kv_heads})." + ) + + if isinstance(config, Trainer.Config): + debug = config.debug + seq_len = config.training.seq_len + max_seq_len = self.max_seq_len + if seq_len > max_seq_len: + raise ValueError( + f"Training sequence length {seq_len} exceeds " + f"attention RoPE maximum supported sequence " + f"length {max_seq_len}." + ) + + for layer_cfg in self.layers: + attention_cfg = getattr(layer_cfg, "attention", None) + if attention_cfg is not None: + attention_cfg.rope = dataclasses.replace( + attention_cfg.rope, max_seq_len=seq_len + ) + if hasattr(layer_cfg, "moe") and layer_cfg.moe is not None: + layer_cfg.moe.router._debug_force_load_balance = ( + debug.moe_force_load_balance + ) + # pyrefly: ignore [bad-override] def get_nparams_and_flops( self, model: nn.Module, seq_len: int diff --git a/torchtitan/models/gpt_oss/sharding.py b/torchtitan/models/gpt_oss/sharding.py index 0732b9a7bc7..1785b577956 100644 --- a/torchtitan/models/gpt_oss/sharding.py +++ b/torchtitan/models/gpt_oss/sharding.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from typing import TYPE_CHECKING @@ -115,10 +116,18 @@ def _set_gpt_oss_layer_sharding( set_qkv_linear_sharding(attention.qkv_linear) attention.wo.sharding_config = scaled_bias_rowwise_config(output_sp=enable_sp) - set_gqa_inner_attention_local_map(attention.inner_attention) + # Inner attentions with their own kernel layout (e.g. the Neuron NKI + # wrapper) set their own local-region placements. + if hasattr(attention.inner_attention, "set_sharding_config"): + attention.inner_attention.set_sharding_config() + else: + set_gqa_inner_attention_local_map(attention.inner_attention) # MoE FFN (all GPT-OSS blocks are MoE). if layer_cfg.moe is not None: + if hasattr(layer_cfg.moe, "set_sharding_config"): + layer_cfg.moe.set_sharding_config(enable_sp=enable_sp, enable_ep=enable_ep) + return set_moe_sharding_config( layer_cfg.moe, enable_ep=enable_ep, diff --git a/torchtitan/models/qwen3/__init__.py b/torchtitan/models/qwen3/__init__.py index 29c337427f8..608f1562dda 100644 --- a/torchtitan/models/qwen3/__init__.py +++ b/torchtitan/models/qwen3/__init__.py @@ -4,7 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. - from collections.abc import Callable from functools import partial diff --git a/torchtitan/models/qwen3/config_registry.py b/torchtitan/models/qwen3/config_registry.py index c2185fd7c0c..b24fa2b217d 100644 --- a/torchtitan/models/qwen3/config_registry.py +++ b/torchtitan/models/qwen3/config_registry.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from torchtitan.components.checkpoint import CheckpointManager from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss @@ -252,6 +253,19 @@ def qwen3_30b_a3b() -> Trainer.Config: ) +def use_global_batch_lbl(config: Trainer.Config) -> Trainer.Config: + """Balance every MoE layer with Qwen3's global-batch load-balancing loss.""" + assert config.model_spec is not None + for layer in config.model_spec.model.layers: + if getattr(layer, "moe", None) is not None: + layer.moe.load_balancing = "global_aux_loss" + return config + + +def qwen3_30b_a3b_global_lbl() -> Trainer.Config: + return use_global_batch_lbl(qwen3_30b_a3b()) + + def qwen3_32b() -> Trainer.Config: model_spec = model_registry("32B") return Trainer.Config( diff --git a/torchtitan/models/qwen3/model.py b/torchtitan/models/qwen3/model.py index 5516ecfc43f..d6aac0ae695 100644 --- a/torchtitan/models/qwen3/model.py +++ b/torchtitan/models/qwen3/model.py @@ -1,10 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. # -# Copyright (c) Meta Platforms, Inc. All Rights Reserved. +# Licensed under the BSD-3-Clause license. from dataclasses import dataclass diff --git a/torchtitan/models/qwen3/sharding.py b/torchtitan/models/qwen3/sharding.py index d93e859fec8..738899652c6 100644 --- a/torchtitan/models/qwen3/sharding.py +++ b/torchtitan/models/qwen3/sharding.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. from typing import TYPE_CHECKING @@ -103,6 +104,14 @@ def _set_qwen3_layer_sharding( # MoE FFN (MoE-enabled layers only). if layer_cfg.moe is not None: + # NKIMoE packs its expert weights along different tensor dimensions than + # standard GroupedExperts, so it must set its own sharding placements. + if hasattr(layer_cfg.moe, "set_sharding_config"): + layer_cfg.moe.set_sharding_config( + enable_sp=enable_sp, + enable_ep=enable_ep, + ) + return set_moe_sharding_config( layer_cfg.moe, enable_ep=enable_ep, diff --git a/torchtitan/protocols/checkpoint_load.py b/torchtitan/protocols/checkpoint_load.py new file mode 100644 index 00000000000..b42bd8b3803 --- /dev/null +++ b/torchtitan/protocols/checkpoint_load.py @@ -0,0 +1,45 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Literal + +import torch + + +_fused_qkv_load_planning: ContextVar[bool] = ContextVar( + "fused_qkv_load_planning", default=False +) + + +@dataclass(frozen=True, slots=True) +class FusedQKVLoadTarget: + """Metadata-only checkpoint target for one projection of a fused QKV tensor.""" + + fused_tensor: torch.Tensor + projection: Literal["q", "k", "v"] + shape: torch.Size + dtype: torch.dtype + n_kv_heads: int + heads_per_kv: int + head_dim: int + + +@contextmanager +def plan_fused_qkv_load() -> Iterator[None]: + """Make FusedQKVLinear state-dict hooks emit metadata-only load targets.""" + + token = _fused_qkv_load_planning.set(True) + try: + yield + finally: + _fused_qkv_load_planning.reset(token) + + +def is_fused_qkv_load_planning() -> bool: + return _fused_qkv_load_planning.get() diff --git a/torchtitan/tools/profiler.py b/torchtitan/tools/profiler.py index 2c337e06d89..36793c56407 100644 --- a/torchtitan/tools/profiler.py +++ b/torchtitan/tools/profiler.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. """Kineto profiler + memory-snapshot lifecycle.""" @@ -18,13 +19,43 @@ from torchtitan.config.function import Function from torchtitan.observability import structured_logger as sl from torchtitan.tools.logging import logger -from torchtitan.tools.utils import device_module +from torchtitan.tools.utils import device_module, device_type # Paths expects by meta internal tooling PROFILE_DIR = "profiling/traces" # Profiler.Config.save_traces_folder default PROFILE_ITER_DIR = "iteration_{step}" # PROFILE_DIR/{PROFILE_ITER_DIR} PROFILE_FILE = "rank{rank}_trace.json.gz" # PROFILE_DIR/PROFILE_ITER_DIR/{PROFILE_FILE} +# --- Neuron (torch_neuronx) profiling ----------------------------------------- +# On Neuron the torch profiler is driven by torch_neuronx: the Neuron runtime is +# exposed as the PrivateUse1 profiler activity and a NeuronConfig is passed as the +# profiler's experimental_config. RUNTIME/CPU_UTIL/HOST_MEMORY (lightweight system +# + host metrics) are always captured; the heavy DEVICE profile (hardware +# instruction traces) reserves ~5GB HBM, so it is opt-in via an env var. +NEURON_PROFILE_DEVICE_ENV = "NEURON_PROFILE_DEVICE" # "1" -> add ProfileMode.DEVICE +# NRT system-trace ring buffer size. The 100K-events/NC NeuronConfig default drops +# the vast majority of events for even one active step of a multi-core distributed +# run (~837K observed); 10M gives headroom. Set via NeuronConfig(max_events_per_nc). +NEURON_MAX_EVENTS_PER_NC = 10_000_000 +# NeuronCores whose system/device trace is captured. Defaults to "0,1" (each .ntff +# is ~1GB, so capturing every core of a 64-core run is huge and the profile is +# usually analyzed rank-0-first); override with NEURON_PROFILE_CAPTURE_NC, e.g. +# "0,1,2,3" or "0-63". +NEURON_CAPTURE_NC_ENV = "NEURON_PROFILE_CAPTURE_NC" +NEURON_DEFAULT_CAPTURE_ENABLED_FOR_NC = "0,1" +# Global-rank allowlist, defaulting to rank 0. This is distinct from +# capture_enabled_for_nc: with one process per NeuronCore, every process sees its +# assigned core as NC0. Override with a comma-separated list, e.g. "0,1,2,3". +NEURON_PROFILE_RANKS_ENV = "NEURON_PROFILE_RANKS" +NEURON_DEFAULT_PROFILE_RANKS = "0" +# On session stop the runtime copies the loaded NEFFs from this dir into the +# session's neffs/, matched to the DEVICE .ntff by the UUID embedded in each NEFF, +# so neuron-explorer can resolve the device trace (else "unable to get NEFFs by +# UID"). torch_neuronx's default compile cache is /tmp/neff_cache; override with +# NEURON_PROFILE_NEFF_CACHE_DIR if the run's NEFF cache lives elsewhere. +NEURON_NEFF_CACHE_DIR_ENV = "NEURON_PROFILE_NEFF_CACHE_DIR" +NEURON_DEFAULT_NEFF_CACHE_DIR = "/tmp/neff_cache" + MEMORY_DIR = ( "profiling/memory_snapshot" # Profiler.Config.save_memory_snapshot_folder default ) @@ -67,6 +98,9 @@ def __init__( self._leaf_folder = leaf_folder self._rank = rank + def set_step_num(self, step_num: int) -> None: + self.step_num = step_num + def step(self, exit_ctx: bool = False) -> None: self.step_num += 1 if not exit_ctx and self.step_num % self.freq != 0: @@ -221,6 +255,7 @@ def active( self._global_step = global_step self._base_folder = base_folder self._leaf_folder = leaf_folder + self.set_memory_profiler_step(global_step) return self def __enter__(self) -> "Profiler": @@ -229,11 +264,7 @@ def __enter__(self) -> "Profiler": base_folder=self._base_folder, leaf_folder=self._leaf_folder, ) - self.memory_profiler = self.build_memory_profiler( - global_step=self._global_step, - base_folder=self._base_folder, - leaf_folder=self._leaf_folder, - ) + self.start_memory_profiler() return self def __exit__(self, exc_type, exc_val, exc_tb) -> bool: @@ -256,6 +287,21 @@ def step(self) -> None: with sl.log_trace_span("memory_profiler_step_call"): self.memory_profiler.step() + def start_memory_profiler(self): + """Start memory-history recording without starting the torch profiler.""" + + if self.memory_profiler is None: + self.memory_profiler = self.build_memory_profiler( + global_step=self._global_step, + base_folder=self._base_folder, + leaf_folder=self._leaf_folder, + ) + return self.memory_profiler + + def set_memory_profiler_step(self, global_step: int) -> None: + if self.memory_profiler is not None: + self.memory_profiler.set_step_num(global_step) + def build_torch_profiler( self, *, @@ -280,6 +326,19 @@ def build_torch_profiler( ) rank = torch.distributed.get_rank() + if device_type == "neuron": + profile_ranks = os.environ.get( + NEURON_PROFILE_RANKS_ENV, NEURON_DEFAULT_PROFILE_RANKS + ) + enabled_ranks = {int(value) for value in profile_ranks.split(",")} + if rank not in enabled_ranks: + logger.info( + "Neuron profiling disabled on rank %d by %s=%s", + rank, + NEURON_PROFILE_RANKS_ENV, + profile_ranks, + ) + return None post_processor = ( cfg.trace_post_processor.build() if cfg.trace_post_processor else None ) @@ -323,23 +382,108 @@ def trace_handler(prof): wait >= 0 ), "profile_freq must be greater than or equal to warmup + active" activities = [torch.profiler.ProfilerActivity.CPU] + profile_kwargs: dict = {} + neuron_trace_handler = None if torch.cuda.is_available(): activities.append(torch.profiler.ProfilerActivity.CUDA) elif torch.xpu.is_available(): activities.append(torch.profiler.ProfilerActivity.XPU) - + elif device_type == "neuron": + neuron_trace_handler = self._enable_neuron_profiling( + activities, profile_kwargs, trace_dir + ) + if neuron_trace_handler is not None and post_processor is not None: + logger.warning( + "trace_post_processor is set but is ignored on Neuron; the " + "Neuron profiler uses its own trace exporter." + ) + + # On Neuron the Neuron exporter fully replaces the chrome-trace handler: + # it writes the native .ntff/.pb session dirs + per-rank framework + # trace.json instead of the chrome PROFILE_FILE, so the chrome + # `trace_handler` and its `trace_post_processor` hook do not run there + # (the warning above surfaces an ignored post_processor). CUDA/CPU/XPU keep + # the chrome `trace_handler`. torch_profiler = torch.profiler.profile( activities=activities, schedule=torch.profiler.schedule( wait=wait, warmup=warmup, active=active, **additional_params ), - on_trace_ready=trace_handler, + on_trace_ready=neuron_trace_handler or trace_handler, record_shapes=True, + **profile_kwargs, ) torch_profiler.__enter__() torch_profiler.step_num = global_step return torch_profiler + def _enable_neuron_profiling( + self, + activities: list, + profile_kwargs: dict, + trace_dir: str, + ): + """Wire torch_neuronx profiling into the torch profiler (in place). + + Appends the Neuron profiler activity to ``activities``, sets the + ``experimental_config``/``with_stack`` entries in ``profile_kwargs``, and + returns the ``NeuronProfiler.export_trace`` handler that writes the + ``.ntff``/``.pb`` session dirs (plus the per-rank framework trace) under + ``trace_dir``. With ``enable_memory_snapshot``, the handler also writes + ``memory_snapshot.pickle`` into the same session dir. + ``RUNTIME``/``CPU_UTIL``/``HOST_MEMORY`` are always enabled; + ``DEVICE`` (hardware instruction traces, ~5GB HBM) is added only when + ``NEURON_PROFILE_DEVICE=1``. + + Returns ``None`` (leaving ``activities``/``profile_kwargs`` untouched) if + ``torch_neuronx`` profiling is unavailable, so the caller falls back to a + CPU-only trace. + """ + try: + from torch_neuronx.profiling import ( + NeuronConfig, + NeuronProfiler, + ProfileMode, + ) + except ImportError: + logger.warning( + "torch_neuronx.profiling unavailable; falling back to CPU-only " + "profiling on Neuron." + ) + return None + + modes = [ProfileMode.RUNTIME, ProfileMode.CPU_UTIL, ProfileMode.HOST_MEMORY] + if os.environ.get(NEURON_PROFILE_DEVICE_ENV, "0") == "1": + modes.append(ProfileMode.DEVICE) + + neff_cache_dir = os.environ.get( + NEURON_NEFF_CACHE_DIR_ENV, NEURON_DEFAULT_NEFF_CACHE_DIR + ) + capture_enabled_for_nc = os.environ.get( + NEURON_CAPTURE_NC_ENV, NEURON_DEFAULT_CAPTURE_ENABLED_FOR_NC + ) + + neuron_config = NeuronConfig( + modes=modes, + profile_output_dir=trace_dir, + max_events_per_nc=NEURON_MAX_EVENTS_PER_NC, + capture_enabled_for_nc=capture_enabled_for_nc, + neff_cache_dir=neff_cache_dir, + ) + activities.append(torch.profiler.ProfilerActivity.PrivateUse1) + profile_kwargs["experimental_config"] = neuron_config + profile_kwargs["with_stack"] = True + logger.info(f"Neuron profiling enabled with modes: {[m.name for m in modes]}") + neuron_profiler = NeuronProfiler(neuron_config) + if not self._config.enable_memory_snapshot: + return neuron_profiler.export_trace + + def export_trace_and_snapshot(prof): + neuron_profiler.export_trace(prof) + neuron_profiler.export_snapshot() + + return export_trace_and_snapshot + def build_memory_profiler( self, *, diff --git a/torchtitan/tools/utils.py b/torchtitan/tools/utils.py index 0deb36d28fa..5d618740fc0 100644 --- a/torchtitan/tools/utils.py +++ b/torchtitan/tools/utils.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import contextlib import gc @@ -40,6 +41,23 @@ def has_rocm_capability(major: int, minor: int) -> bool: ) +def has_neuron_capability(feature: str) -> bool: + """Whether the current Neuron platform supports ``feature``. + + Currently only ``"mxfp8"`` is a real capability, which requires trn3 + (NeuronCore-v4) or later. Returns False off-Neuron. + """ + if feature.lower() != "mxfp8": + logger.warning(f"Unknown neuron capability {feature} not supported") + return False + try: + import torch_neuronx + + return torch_neuronx.utils.get_platform_target().startswith("trn3") + except Exception: + return False + + def get_device_info() -> tuple[str, ModuleType]: device_type = _get_available_device_type() or "cuda" device_module = _get_device_module(device_type) # default device_module:torch.cuda @@ -47,6 +65,13 @@ def get_device_info() -> tuple[str, ModuleType]: device_type, device_module = get_device_info() +if device_type == "neuron": + # FlexAttention binds torch.compile into a class variable when its module is + # imported, so the Neuron default backend must be installed before model imports. + # TODO: Replace this import-time hook after the upstream design is approved. + from torchtitan_neuron.compile_backend import install_neuron_compile_backend + + install_neuron_compile_backend() # used to avoid stragglers in garbage collection diff --git a/torchtitan/trainer.py b/torchtitan/trainer.py index 1c4b3f82f99..361b8796ef2 100644 --- a/torchtitan/trainer.py +++ b/torchtitan/trainer.py @@ -1,8 +1,9 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. +# Modifications Copyright (c) Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the BSD-3-Clause license. import dataclasses import json @@ -44,10 +45,17 @@ MemoryBudgetAC, SelectiveAC, ) -from torchtitan.distributed.context_parallel import prepare_context_parallel_input +from torchtitan.distributed.context_parallel import ( + apply_cp_to_rope, + prepare_context_parallel_input, +) from torchtitan.distributed.spmd_types import annotate_input_spmd_types from torchtitan.models.common.attention import FlexAttention, VarlenAttention from torchtitan.models.common.decoder import Decoder +from torchtitan.models.common.moe import ( + configure_global_batch_lbl, + uses_global_batch_lbl, +) from torchtitan.observability import structured_logger as sl from torchtitan.protocols import BaseModel from torchtitan.protocols.model_spec import ModelSpec @@ -210,7 +218,9 @@ def maybe_log(self) -> None: # runtime utilities device: torch.device gc_handler: utils.GarbageCollection + profiler: Profiler train_context: dist_utils.SpmdContext + compiled_grad_clip: dist_utils.CompiledGradClip | None = None gradient_accumulation_steps: int pp_has_first_stage: bool pp_has_last_stage: bool @@ -235,6 +245,7 @@ def __init__(self, config: Config): self.device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}") # Device has to be set before creating TorchFT manager. device_module.set_device(self.device) + self.compiled_grad_clip = dist_utils.build_compiled_grad_clip(config.compile) # init distributed and build meshes self.parallel_dims = parallel_dims = self.init_distributed() @@ -301,6 +312,13 @@ def __init__(self, config: Config): ): model = model_config.build() + if parallel_dims.cp_enabled: + apply_cp_to_rope( + model, + parallel_dims.get_mesh("cp"), + load_balancer_type=config.parallelism.context_parallel_load_balancer, + ) + # Verify all submodules satisfy the Module protocol # TODO: move this to module validate(). # This is current put here to verify module build and @@ -367,6 +385,19 @@ def __init__(self, config: Config): ) assert self.gradient_accumulation_steps > 0 + if parallel_dims.spmd_backend != "default" and uses_global_batch_lbl(model): + # full_dtensor would double-reduce DP/CP counts; spmd_types is not validated for this loss. + raise ValueError( + "load_balancing='global_aux_loss' needs spmd_backend='default'" + ) + + self.profiler = config.profiler.build( + global_step=0, + base_folder=config.dump_folder, + ) + if config.profiler.enable_memory_snapshot: + self.profiler.start_memory_profiler() + # apply parallelisms and initialization with sl.log_trace_span("model_parallelism_init"): if parallel_dims.pp_enabled: @@ -436,6 +467,34 @@ def __init__(self, config: Config): self.model_parts = [model] + # Global-batch load balancing: every MoE counts tokens over the ranks holding distinct tokens + # and shares the step's normalization over all its forwards (gradient accumulation x PP + # micro-batches), as the main loss does through global_valid_tokens. + num_forwards = self.gradient_accumulation_steps * ( + config.training.local_batch_size + // config.parallelism.pipeline_parallel_microbatch_size + if parallel_dims.pp_enabled + else 1 + ) + configure_global_batch_lbl( + self.model_parts, + batch_mesh=parallel_dims.get_optional_mesh("loss"), + num_forwards=num_forwards, + ) + + # Modules carrying weight-derived state that must be rebuilt after + # init_weights, after checkpoint load, and after each optimizer step (e.g. + # MXFP8 pre-quantized SimpleFSDP weights). Duck-typed so core stays + # device-agnostic; the model config decides whether any such module exists. + # Must prime after init_weights (shards hold real data), before first forward. + self._weight_derived_state_modules = [ + module + for part in self.model_parts + for module in part.modules() + if hasattr(module, "refresh_weight_derived_state") + ] + self._refresh_weight_derived_state() + # Set lm_head reference for ChunkedLossWrapper after model construction. # Non-PP: single model part always has lm_head. # PP: only the last stage has lm_head; non-last stages skip this. @@ -755,6 +814,11 @@ def forward_backward_step( # The returned loss here is local SUM loss / global_valid_tokens return loss + def _refresh_weight_derived_state(self) -> None: + """Rebuild any per-module weight-derived state (no-op when none exists).""" + for module in getattr(self, "_weight_derived_state_modules", ()): + module.refresh_weight_derived_state() + def train_step( self, data_iterator: Iterator[tuple[dict[str, torch.Tensor], torch.Tensor]] ): @@ -809,11 +873,16 @@ def train_step( foreach=True, pp_mesh=parallel_dims.get_optional_mesh("pp"), ep_enabled=parallel_dims.ep_enabled, + compiled_grad_clip=self.compiled_grad_clip, ) self.checkpointer.maybe_wait_for_staging() self.optimizers.step() self.lr_schedulers.step() + # Not an optimizer post-hook: reaching for the full weight there would + # all-gather bf16 on a schedule the compiled graph does not expect. + self._refresh_weight_derived_state() + # Reduce the data collected over gradient accumulation steps. loss = torch.sum(torch.stack(accumulated_losses)) @@ -881,15 +950,23 @@ def train(self): sl.log_trace_instant("training_start") + profiler = self.profiler + profiler.active(global_step=self.step, base_folder=config.dump_folder) + self.checkpointer.load(step=config.checkpoint.load_step) + # Loading replaces the weights the __init__ priming described. Without this, + # step 1 alone runs on pre-checkpoint encodings. + self._refresh_weight_derived_state() + # Capture loaded step for relative_step calculation. # After checkpoint load: self.step = restored step (e.g. 100), or 0 if fresh. loaded_step = self.step + profiler.set_memory_profiler_step(loaded_step) logger.info(f"Training starts at step {self.step + 1}") - with config.profiler.build( + with profiler.active( global_step=self.step, base_folder=config.dump_folder, ) as profiler: diff --git a/torchtitan_neuron/README.md b/torchtitan_neuron/README.md new file mode 100644 index 00000000000..9a18dbcf20f --- /dev/null +++ b/torchtitan_neuron/README.md @@ -0,0 +1,66 @@ +# TorchTitan on AWS Neuron + +`torchtitan_neuron` contains the Neuron-specific components used to train +TorchTitan models on AWS Trainium. It is installed alongside `torchtitan` and +uses the same configuration and launch interfaces. + +## Package boundary + +This package owns self-contained Neuron kernels, kernel wrappers, recipes, +configuration builders, tests, and documentation. It may import the +device-independent TorchTitan core. Core code should only import +`torchtitan_neuron` lazily at a routing boundary. + +## Layout + +``` +torchtitan_neuron/ +├── models/ Neuron modules and Trainer config registries. +├── experiments/ GraphTrainer config registries and shared conversion helpers. +├── docs/ Neuron-specific documentation. +└── tests/ Neuron kernel, Trainer, and GraphTrainer tests. +``` + +## Included components + +- DeepSeek-V3 NKI MoE Trainer and GraphTrainer recipes. +- GPT-OSS NKI attention and MoE Trainer and GraphTrainer recipes. +- Llama3 dense Trainer and SimpleFSDP GraphTrainer recipes. +- Qwen3 dense/MXFP8 Trainer and NKI MoE Trainer/GraphTrainer recipes. +- The shared `neuron_simple_fsdp` GraphTrainer converter. + +Select a recipe by passing its registry module and configuration name to the +TorchTitan training entrypoint. + +For example, the Qwen3 dense prefix-cache recipe is selected explicitly: + +```bash +NGPU=64 MODULE=torchtitan_neuron.models.qwen3 \ + CONFIG=qwen3_8b_tp4_fsdp16_sac ./run_train.sh +``` + +RoPE remains a core component. Recipes that select sequential prefix positions +set `layer.attention.rope.use_prefix_cache_for_positions = True`; importing +`torchtitan_neuron` or selecting a device does not enable this behavior. +This preserves the earlier overrides' position handling, including ignoring +document position resets. Core recipes retain their default position gather. + + +## Attention + +Llama3, DeepSeek-V3, and Qwen3 Neuron recipes use causal SDPA attention. +GPT-OSS uses NKI attention with native attention sinks and alternates between +128-token causal windows and full causal attention. The recipe selects attention +independently of its MoE implementation and compilation backend. + +These attention paths do not isolate documents packed into the same sequence. +See [Neuron attention backends](docs/neuron_attention_backends.md) for masking +behavior, recipe selection, launch instructions, and training configuration. + +## Documentation + +- [Dense model tutorials](docs/dense-tutorials.md) +- [MoE model tutorials](docs/moe-tutorials.md) +- [Neuron attention backends](docs/neuron_attention_backends.md) +- [Qwen3 NKI MoE](docs/qwen3_nki_moe.md) +- [Full Qwen3-235B training on a Trn3 UltraServer](docs/qwen3_235b_full_ep128.md) diff --git a/torchtitan_neuron/__init__.py b/torchtitan_neuron/__init__.py new file mode 100644 index 00000000000..3246efdf889 --- /dev/null +++ b/torchtitan_neuron/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""torchtitan_neuron: Neuron (AWS Trainium) specific extensions for torchtitan. + +Holds Neuron-specific kernels, ops, and configs (e.g. NKI MoE) kept out of +core torchtitan. Neuron recipes select them by converting the stock model +config (for example ``to_nki_moe_config`` for each model's NKI MoE). +""" diff --git a/torchtitan_neuron/compile_backend.py b/torchtitan_neuron/compile_backend.py new file mode 100644 index 00000000000..d66e3ba6c11 --- /dev/null +++ b/torchtitan_neuron/compile_backend.py @@ -0,0 +1,40 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +import functools +from typing import Any, cast + +import torch + + +def install_neuron_compile_backend() -> None: + """Route bare ``torch.compile`` calls to the Neuron dynamo backend. + + Several call sites (flex attention, block-mask compilation) call + ``torch.compile`` with no backend, so they default to inductor, which + cannot lower to the neuron device ("No interface for device neuron"). + Rather than thread a backend through every site, wrap ``torch.compile`` + once: on the neuron device, and only when the caller did not pass a + ``backend``, use the registered "neuron" backend. Any ``options`` the caller + passed are left intact -- the neuron backend accepts them, and the flex call + sites pass their inductor autotune configs there. Calls that pass an explicit + ``backend`` (e.g. via ``config.compile.backend``) are untouched. + + The caller is responsible for invoking this only on the Neuron device. + """ + if getattr(torch.compile, "_neuron_wrapped", False): + return + + _original_compile = torch.compile + + @functools.wraps(_original_compile) + def _compile(model=None, **kwargs): + if "backend" not in kwargs: + kwargs["backend"] = "neuron" + return _original_compile(model, **kwargs) + + compile_wrapper = cast(Any, _compile) + compile_wrapper._neuron_wrapped = True + torch.compile = compile_wrapper diff --git a/torchtitan_neuron/components/__init__.py b/torchtitan_neuron/components/__init__.py new file mode 100644 index 00000000000..fcb62242c2d --- /dev/null +++ b/torchtitan_neuron/components/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. diff --git a/torchtitan_neuron/components/quantization/__init__.py b/torchtitan_neuron/components/quantization/__init__.py new file mode 100644 index 00000000000..771029f4bc1 --- /dev/null +++ b/torchtitan_neuron/components/quantization/__init__.py @@ -0,0 +1,33 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron-only quantization extensions kept out of core torchtitan. + +Holds the MXFP8 pre-quantized SimpleFSDP weight path (``mx_prequant``), which +depends on torchao's ``neuron::`` custom ops. Opt in per recipe with +``NeuronMXFP8LinearConverter(prequantize_weights=True)``; the trainer refreshes +the encodings via the generic ``refresh_weight_derived_state`` hook. +""" + +import importlib.util + + +def torchao_supports_neuron() -> bool: + """True iff torchao is installed and exposes Neuron MXFP8 support. + + The Neuron MXFP8 path needs the torchao build that adds + ``KernelPreference.NEURON`` and the ``neuron::`` custom ops. Public torchao (a + plain ``pip install torchao``) lacks it; callers skip or raise rather than + crash. The fork is not vended publicly yet. + """ + if importlib.util.find_spec("torchao") is None: + return False + try: + from torchao.quantization.quantize_.common.kernel_preference import ( + KernelPreference, + ) + except Exception: + return False + return hasattr(KernelPreference, "NEURON") diff --git a/torchtitan_neuron/components/quantization/mx_prequant.py b/torchtitan_neuron/components/quantization/mx_prequant.py new file mode 100644 index 00000000000..3c788552c5a --- /dev/null +++ b/torchtitan_neuron/components/quantization/mx_prequant.py @@ -0,0 +1,544 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Post-optimizer MXFP8 weight quantization for SimpleFSDP. + +The stock MXFP8 path all-gathers a bf16 weight and quantizes it inside the GEMM every +step. This path quantizes each shard once after ``optimizer.step()``, so the forward +all-gathers 1-byte fp8 plus its e8m0 scale instead of bf16. + +Opt in through the config: attach ``NeuronMXFP8LinearConverter`` with +``prequantize_weights=True`` in a recipe. The converter swaps each ``MXFP8Linear`` for a +``NeuronMXFP8Linear`` whose ``forward`` consumes pre-quantized encodings and whose +``refresh_weight_derived_state`` re-quantizes the local shard. The trainer calls that +method (duck-typed) after init_weights, after checkpoint load, and after each optimizer +step -- the analogue of ``precompute_float8_dynamic_scale_for_fsdp``. + +Two encodings are needed because MX block scaling runs along the *contraction* axis, +which differs between the two GEMMs that consume the weight: + + forward out = x @ W.T contracts in_features -> blocks along in + dgrad grad_x = grad_out @ W contracts out_features -> blocks along out + +so the fprop encoding comes from ``W`` and the dgrad one from ``W.t()``, both blocked on +the last dim. The transpose has to happen eagerly, here: an in-graph ``.t()`` on a +pre-quantized operand loses the packing the compiler requires. + +Both encodings live on ``_PrequantWeightOperands``; see it for their shapes. Each shard +dim plays both roles across the two encodings, so both dims must satisfy both kernel +bounds -- see ``MIN_K_BY_SCHEME`` / ``MIN_F``, enforced by ``_check_shard`` on the first +refresh. Both bounds are met only by the dim-1-FSDP-sharded dev flavors, which require +SimpleFSDP dim-1 sharding support. + +``prequant_scheme`` picks the packed layout: ``wrapX`` (nkilib's own quantize kernel) or +``1x32``. They differ only in which K occupies which packed slot, so the same value must +reach the quantize kernel and the matmul. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, fields + +import torch +from torch.distributed.tensor import DTensor, Partial, Replicate, Shard +from torch.distributed.tensor.placement_types import _StridedShard + +from torchtitan.components.quantization.mx import MXFP8Linear, MXFP8LinearConverter + +logger = logging.getLogger(__name__) + +FP8_DTYPE = torch.float8_e4m3fn + +# nkilib QuantScheme values, spelled as matmul_mxfp8's ``quant_scheme`` expects them. +SCHEME_WRAPX = "wrapX" +SCHEME_1X32 = "1x32" +_SCHEMES = (SCHEME_WRAPX, SCHEME_1X32) + +# nkilib quantize_block_mxfp8_kernel bounds: MIN_K_FOR_QUANTIZATION (128) on the +# contraction dim, MIN_F_FOR_QUANTIZATION on the free dim. +# +# The floor on K is the KERNEL's 128 under 1x32 and 512 under wrapX, and the difference +# is about ordering, not the kernel. A gathered encoding is per-rank encodings +# concatenated along packed dim 0. Under 1x32 that axis is K-consecutive, so any chunk +# boundary is a K boundary and only the packing's /4 and the 32-element block matter. +# Under wrapX the kernel interleaves each 512-element K tile with stride 128, so per-rank +# tilings coincide with the full-K tiling only when each chunk is a whole 512-tile; a +# sub-512 chunk uses strides that never occur in the full-K decomposition and the matmul +# then pairs mismatched K silently (~100% error, no exception). +MIN_K_BY_SCHEME = {SCHEME_WRAPX: 512, SCHEME_1X32: 128} +MIN_F = 8 + +__all__ = ["NeuronMXFP8Linear", "NeuronMXFP8LinearConverter"] + +# Instance attributes the refresh writes and the forward reads. +_OPERANDS = "_mxfp8_operands" +_SPEC = "_mxfp8_mesh_spec" + + +def _canonical_scheme(scheme: str) -> str: + """Validate a scheme name and return its canonical spelling. + + Spelling matters: it is passed through to nkilib, which compares against its own + QuantScheme values, so accept case-insensitively but hand on the canonical form. + """ + for known in _SCHEMES: + if scheme.lower() == known.lower(): + return known + raise ValueError( + f"prequant_scheme={scheme!r} is not a known scheme; expected one of {_SCHEMES}" + ) + + +@dataclass(frozen=True, slots=True) +class _PrequantWeightOperands: + """The packed MXFP8 encodings derived from one local weight shard. + + Named and shaped after upstream's ``_MXFP8LinearOperands`` + (``torchtitan/components/quantization/mxfp8/tensor.py``, pytorch/torchtitan#4203) so + adopting that overhaul is a rename rather than a rewrite. It cannot share upstream's + field *set*: there, square 32x32 scale tiles are transpose-invariant, so one qdata + allocation serves both GEMMs and only the scales differ. Neuron's quantize kernels + offer no 2D scheme -- nkilib ``QuantScheme`` is ``WRAPX | _1x32``, both blocked along + K alone -- so fprop and dgrad are independently quantized and each owns its qdata. + + For a local shard ``[O, I]``, contraction dim leading the packed layout: + + fprop (contracts I) qdata [I/4, O*4] scale [I/4, O] + dgrad (contracts O) qdata [O/4, I*4] scale [O/4, I] + + Fields are distinct allocations, which is what lets the refresh refill them in + place. Field order is the gather order in ``_prequant_forward``. + """ + + qdata_fprop: torch.Tensor + scale_fprop: torch.Tensor + qdata_dgrad: torch.Tensor + scale_dgrad: torch.Tensor + + def tensors(self) -> tuple[torch.Tensor, ...]: + return tuple(getattr(self, f.name) for f in fields(self)) + + +def _quantize_mx(t: torch.Tensor, scheme: str): + """Quantize to packed MXFP8, contracting over ``t``'s LAST dim. + + Returns ``(data [K/4, F*4], scales [K/4, F])`` for ``t`` of shape ``[F, K]`` -- the + swizzled x4-packed layout ``matmul_mxfp8`` requires of a pre-quantized operand. Note + the contraction dim leads the packed result, which is why the gather placements below + invert relative to the weight's own dims. + """ + if t.dtype not in (torch.bfloat16, torch.float16): + t = t.to(torch.bfloat16) + if not t.is_contiguous(): + t = t.contiguous() + return torch.ops.neuron.quantize_mx_nki(t, FP8_DTYPE, scheme) + + +def _build_operands( + shard: torch.Tensor, + scheme: str, + out: _PrequantWeightOperands | None = None, +) -> _PrequantWeightOperands: + """Quantize a local weight shard into both encodings. + + Allocates a new operands when ``out`` is None, otherwise refills ``out``'s existing + tensors in place and returns it. Mirrors upstream's ``_build_operands`` contract. + + The dgrad encoding materialises ``[I, O]`` here, eagerly: an in-graph ``.t()`` on a + pre-quantized operand loses the packing the compiler requires. + """ + fresh = _PrequantWeightOperands( + *_quantize_mx(shard, scheme), *_quantize_mx(shard.t(), scheme) + ) + if out is None: + return fresh + for existing, value in zip(out.tensors(), fresh.tensors(), strict=True): + # A raise, not an assert: this guards against silent corruption rather than a + # programmer typo, and asserts vanish under -O. copy_ broadcasts, so a mismatch + # slipping through would write a wrong-but-plausible encoding instead of failing. + if existing.shape != value.shape or existing.dtype != value.dtype: + raise RuntimeError( + f"encoding changed shape/dtype across refresh: " + f"{tuple(existing.shape)}/{existing.dtype} -> " + f"{tuple(value.shape)}/{value.dtype}" + ) + existing.copy_(value) + return out + + +def _check_shard(name: str, o: int, i: int, scheme: str) -> None: + """Validate a local weight shard against the NKI quantize kernel's constraints. + + The kernel needs its contraction dim (``K``) a multiple of the active scheme's + ``MIN_K`` and its free dim (``F``) a multiple of ``MIN_F``. Both encodings are taken + from the same shard with the roles swapped, so each dim has to satisfy both bounds. + """ + min_k = MIN_K_BY_SCHEME[scheme] + for dim_name, k, f in (("fprop", i, o), ("dgrad", o, i)): + if k % min_k: + raise ValueError( + f"{name}: {dim_name} encoding contracts over {k}, which must be a " + f"multiple of {min_k} under scheme {scheme}; got {k} % {min_k} = " + f"{k % min_k}. Either a block straddles two ranks (so the local amax is " + f"not the block amax) or the per-rank K tiling does not line up with the " + f"full-K tiling. Shard the other axis, reduce the FSDP degree, widen the " + f"model dim, or -- if on wrapX -- try the 1x32 scheme, whose floor is " + f"{MIN_K_BY_SCHEME[SCHEME_1X32]}." + ) + if f % MIN_F: + raise ValueError( + f"{name}: {dim_name} encoding's free dim ({f}) must be a multiple of " + f"{MIN_F}; got {f} % {MIN_F} = {f % MIN_F}." + ) + + +class _PrequantMM(torch.autograd.Function): + """``x @ W.T`` from pre-quantized weight encodings. + + Only the weight arrives pre-quantized; activations change every step, so ``fprop`` and + ``dgrad`` quantize theirs in-kernel and ``wgrad`` quantizes both operands inline. + + ``weight`` is passed but never read -- the encodings carry the values. It exists + purely as an autograd edge back to the parameter, since the encodings are plain fp8 + tensors with no grad path; without it ``grad_w`` is silently discarded and every + linear stays frozen. Must be the RAW sharded parameter, not ``mod.weight``: that + property runs the SimpleFSDP parametrization and re-introduces the bf16 all-gather + this path exists to avoid. + """ + + @staticmethod + def forward( + ctx, + x, + weight, + qdata_fprop, + scale_fprop, + qdata_dgrad, + scale_dgrad, + wspec, + scheme, + ): + out = torch.ops.neuron.prequantized_mx_scaled_mm_nki( + x, qdata_fprop, scale_fprop, x.dtype, FP8_DTYPE, scheme + ) + ctx.scheme = scheme + ctx.save_for_backward(x, qdata_dgrad, scale_dgrad) + ctx.wspec = wspec + return out + + @staticmethod + def backward(ctx, grad_out): + x, qdata_dgrad, scale_dgrad = ctx.saved_tensors + grad_out = grad_out.contiguous() + + # grad_x = grad_out @ W : contract out_features, so rhs is the [I, O] encoding. + grad_x = torch.ops.neuron.prequantized_mx_scaled_mm_nki( + grad_out, + qdata_dgrad, + scale_dgrad, + grad_out.dtype, + FP8_DTYPE, + ctx.scheme, + ) + + # grad_W = grad_out.T @ x : both operands are activations, so the fully fused op + # quantizes them inline. It computes lhs @ rhs.T, so rhs is x.T. + grad_w = torch.ops.neuron.fused_quantize_mx_scaled_mm_nki( + grad_out.t().contiguous(), + x.t().contiguous(), + grad_out.dtype, + FP8_DTYPE, + ) + # Positions: x, weight, qdata_fprop, scale_fprop, qdata_dgrad, scale_dgrad, + # wspec, scheme + return ( + grad_x, + _to_param_layout(grad_w, ctx.wspec), + None, + None, + None, + None, + None, + None, + ) + + +def _to_param_layout(grad_w, wspec): + """Reshape ``grad_w`` from the activation layout into the parameter's own sharding. + + Each DP rank saw a different batch slice, so its ``grad_w`` is one addend of the true + gradient. Declaring the DP axis ``Partial`` and redistributing to the parameter's + placements is exactly the reduce-scatter SimpleFSDP's parametrization backward would + have done, so this costs no more than the non-prequant path. + + TP placements are read off ``grad_w``, not copied from the parameter: DTensor may have + satisfied the GEMM by redistributing ``x`` to ``Replicate`` on TP, leaving ``grad_w`` + TP-whole, and asserting the parameter's ``Shard`` would then overstate the global + shape by the TP degree (autograd rejects the gradient's shape). + """ + if wspec is None: + return grad_w + mesh, dp_i, raw_placements = wspec + if isinstance(grad_w, DTensor): + local = grad_w.to_local() + tp_placements = list(grad_w.placements) + else: + local = grad_w + tp_placements = [Replicate()] * (mesh.ndim - 1) + + src, tp_at = [], 0 + for i in range(mesh.ndim): + if i == dp_i: + src.append(Partial()) + else: + src.append( + tp_placements[tp_at] if tp_at < len(tp_placements) else Replicate() + ) + tp_at += 1 + + full = DTensor.from_local(local, mesh, src, run_check=False) + return full.redistribute(placements=list(raw_placements)) + + +_DP_MESH_NAMES = ("dp_shard_cp", "dp_shard", "fsdp", "dp") + + +def _retarget(p, new_dim): + """Same placement type, rebased onto ``new_dim``. + + ``_StridedShard``'s ``split_factor`` must be carried over: when TP has already + split a dim, FSDP's split of that same dim is interleaved, and reassembling it as + a plain ``Shard`` would concatenate the pieces in the wrong order -- silently + wrong weights rather than an error. + """ + if isinstance(p, _StridedShard): + return _StridedShard(new_dim, split_factor=p.split_factor) + if isinstance(p, Shard): + return Shard(new_dim) + return p + + +def _placements(raw_placements, transposed: bool): + out = [] + for p in raw_placements: + d = getattr(p, "dim", None) + if d in (0, 1) and transposed: + out.append(_retarget(p, 1 - d)) + else: + out.append(p) + return list(out) + + +def _build_spec(raw): + """Gather plan for one weight, from the RAW parameter (not the property). + + ``mod.weight`` runs SimpleFSDP's parametrization and returns the already-gathered + tensor on the TP-only mesh. ``mod._parameters["weight"]`` is the only place the + ``fsdp`` mesh axis and its ``_StridedShard`` placement are visible. + + Placements are in PACKED coordinates, contraction dim leading: the fprop encoding + contracts in_features, so weight dim 1 becomes packed dim 0 (the ``transposed`` + mapping); the dgrad encoding contracts out_features, so dims carry over unchanged. + Chunk sizes stay integral under the packing's /4 and *4 because ``_check_shard`` + already required each local dim to be a multiple of the scheme's ``MIN_K`` and + ``MIN_F``. + + Returns (mesh, dp_index, fprop_placements, dgrad_placements, tp_mesh, tp_fprop, tp_dgrad) + or None when there is no DP axis to gather. + """ + if not isinstance(raw, DTensor): + return None + mesh = raw.device_mesh + names = list(mesh.mesh_dim_names or []) + dp_i = next((i for i, n in enumerate(names) if n in _DP_MESH_NAMES), None) + if dp_i is None: + return None + tp_names = tuple(n for i, n in enumerate(names) if i != dp_i) + tp_mesh = mesh[tp_names] if tp_names else None + tp_idx = [i for i in range(len(names)) if i != dp_i] + tp_raw = [raw.placements[i] for i in tp_idx] + return ( + mesh, + dp_i, + _placements(raw.placements, True), + _placements(raw.placements, False), + tp_mesh, + _placements(tp_raw, True), + _placements(tp_raw, False), + ) + + +def _gather(local, spec, transposed: bool): + """All-gather one local encoding over the DP axis only, inside the traced forward. + + Mirrors ``ReplicateComputation.replicate_compute``: redistribute the DP axis to + Replicate leaving TP alone, then re-wrap on the TP mesh so the result is a DTensor on + the same mesh as the activation (a plain tensor would mix Tensor and DTensor and raise). + + fp8 is not a supported collective dtype, so the payload is sent as ``uint8``. + Both float8_e4m3fn and float8_e8m0fnu are 1 byte, so the view is + shape-preserving and the placements stay valid. + """ + if spec is None: + return local + mesh, dp_i, pl_fprop, pl_dgrad, tp_mesh, tp_fprop, tp_dgrad = spec + pl = list(pl_dgrad if transposed else pl_fprop) + + # Capture the true element dtype BEFORE the byte view: the receiver has to know + # what the bytes were, and guessing wrong is silently wrong rather than an error. + elem_dtype = local.dtype + payload = local if elem_dtype == torch.uint8 else local.view(torch.uint8) + + d = DTensor.from_local(payload, mesh, pl, run_check=False) + target = list(pl) + # DTensor issues all gather on DP axis and no other + target[dp_i] = Replicate() + gathered = d.redistribute(placements=target).to_local() + if elem_dtype != torch.uint8: + gathered = gathered.view(elem_dtype) + if tp_mesh is None: + return gathered + return DTensor.from_local( + gathered, tp_mesh, tp_dgrad if transposed else tp_fprop, run_check=False + ) + + +def _prequant_forward(self, input: torch.Tensor) -> torch.Tensor: + """``MXFP8Linear.forward`` that consumes pre-quantized weights.""" + operands = getattr(self, _OPERANDS, None) + if operands is None: + raise RuntimeError( + "MXFP8 pre-quantized weights are missing. refresh_weight_derived_state() " + "must run once after init_weights and before the first forward." + ) + orig_shape = input.shape + x = input.reshape(-1, orig_shape[-1]) + spec = getattr(self, _SPEC) + # The RAW sharded parameter, not self.weight: reading the property would run the + # parametrization and all-gather the bf16 weight. Passed only as a grad edge. + raw = self._parameters.get("weight") + wspec = None + if spec is not None and isinstance(raw, DTensor): + wspec = (spec[0], spec[1], tuple(raw.placements)) + + # The encodings cross the autograd boundary as separate positional tensors, not as + # the operands object: tensors reached through a dataclass argument are not graph + # inputs to a traced autograd.Function, so the collectives producing them would fall + # outside the graph. The grouping is for the module state and the quantizer. + out = _PrequantMM.apply( + x, + raw, + _gather(operands.qdata_fprop, spec, False), + _gather(operands.scale_fprop, spec, False), + _gather(operands.qdata_dgrad, spec, True), + _gather(operands.scale_dgrad, spec, True), + wspec, + self._prequant_scheme, + ) + out = out.reshape(*orig_shape[:-1], out.shape[-1]) + if self.bias is not None: + out = out + self.bias.to(out.dtype) + return out + + +if MXFP8Linear is not None: + + class NeuronMXFP8Linear(MXFP8Linear): + """MXFP8Linear that pre-quantizes its weight shard once per optimizer step. + + ``forward`` consumes the packed encodings; ``refresh_weight_derived_state`` + rebuilds them from the current shard and is called by the trainer after + init_weights, after checkpoint load, and after each optimizer step. + """ + + @dataclass(kw_only=True, slots=True) + class Config(MXFP8Linear.Config): + prequant_scheme: str = SCHEME_WRAPX + + def __init__(self, config: Config): + super().__init__(config) + self._prequant_scheme = _canonical_scheme(config.prequant_scheme) + setattr(self, _OPERANDS, None) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + return _prequant_forward(self, input) + + @torch.no_grad() + def refresh_weight_derived_state(self) -> None: + """Re-quantize this rank's weight shard (fprop + dgrad encodings). + + Reads ``_parameters["weight"]``, NOT ``self.weight``: the latter is a property + that runs SimpleFSDP's parametrization and returns the already-gathered weight. + We want the local shard -- quantize 1/FSDP of the weight and all-gather 1 byte + per element instead of 2. + """ + raw = self._parameters.get("weight") + local = raw._local_tensor if isinstance(raw, DTensor) else raw + if local is None or local.ndim != 2: + return + inner = getattr(local, "_tensor", local) + + # Mesh layout is fixed for the run, so resolve the gather plan once, and + # validate the local shard against the kernel bounds on that first pass. + if not hasattr(self, _SPEC): + _check_shard( + type(self).__name__, + local.shape[0], + local.shape[1], + self._prequant_scheme, + ) + setattr(self, _SPEC, _build_spec(raw)) + + # Refill IN PLACE. Rebinding would hand the module new tensor objects each step + # while the compiled graph still holds the ones it captured on the first trace, + # so the linears would keep reading step-0 encodings. No version-counter guard + # is needed: refresh runs after backward has consumed them. + operands = getattr(self, _OPERANDS) + refilled = _build_operands(inner, self._prequant_scheme, out=operands) + if operands is None: + setattr(self, _OPERANDS, refilled) + + class NeuronMXFP8LinearConverter(MXFP8LinearConverter): + """MXFP8 Linear converter with an opt-in pre-quantized SimpleFSDP weight path.""" + + @dataclass(kw_only=True, slots=True) + class Config(MXFP8LinearConverter.Config): + prequantize_weights: bool = False + prequant_scheme: str = SCHEME_WRAPX + + def __init__(self, config: Config): + super().__init__(config) + if config.prequantize_weights: + _canonical_scheme(config.prequant_scheme) + if config.kernel_preference != "neuron": + raise ValueError( + "prequantize_weights=True requires kernel_preference='neuron'" + ) + + def convert(self, model_config): + model_config = super().convert(model_config) + if not self.config.prequantize_weights: + return model_config + for _fqn, cfg, parent, attr in model_config.traverse(MXFP8Linear.Config): + new_config = NeuronMXFP8Linear.Config( + **{f.name: getattr(cfg, f.name) for f in fields(cfg)}, + prequant_scheme=self.config.prequant_scheme, + ) + if parent is None: + model_config = new_config + elif isinstance(parent, list): + parent[attr] = new_config + else: + setattr(parent, attr, new_config) + logger.info( + "MXFP8 pre-quantized SimpleFSDP weight path enabled (scheme=%s)", + self.config.prequant_scheme, + ) + return model_config + +else: + NeuronMXFP8Linear = None + NeuronMXFP8LinearConverter = None diff --git a/torchtitan_neuron/docs/dense-tutorials.md b/torchtitan_neuron/docs/dense-tutorials.md new file mode 100644 index 00000000000..d595ece3d2a --- /dev/null +++ b/torchtitan_neuron/docs/dense-tutorials.md @@ -0,0 +1,199 @@ +# Training dense models with torchtitan-neuron + +## Setup + +```shell +cd /workspace/torchtitan_neuron +pip install -r requirements.txt +pip install -e . +``` + +## Download the Qwen3-32B tokenizer + +```shell +python scripts/download_hf_assets.py --repo_id Qwen/Qwen3-32B --assets tokenizer +``` + +## Train Qwen3-32B (Eager Mode) + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.models.qwen3 \ + --config qwen3_32b_tp4_fsdp16_fullac_eager \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --training.local_batch_size=1 \ + --training.global_batch_size=16 +``` + +## Train Qwen3-32B with torch.compile + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.models.qwen3 \ + --config qwen3_32b_tp4_fsdp16_fullac_compile \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --training.local_batch_size=1 \ + --training.global_batch_size=16 +``` + +## Train Qwen3-32B with SimpleFSDP + torch.compile + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.experiments.graph_trainer.qwen3 \ + --config simple_fsdp_qwen3_32b_tp4_fsdp16_fullac \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --parallelism.context_parallel_degree=1 \ + --training.local_batch_size=1 \ + --training.global_batch_size=16 +``` + +## Download the Llama3-8B tokenizer + +```shell +python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=... +``` + +## Train Llama3-8B (Eager Mode) + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.models.llama3 \ + --config llama3_8b_sac_eager \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --parallelism.context_parallel_degree=1 \ + --training.global_batch_size=16 +``` + +## Train Llama3-8B with torch.compile + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.models.llama3 \ + --config llama3_8b_sac_compile \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --parallelism.context_parallel_degree=1 \ + --training.global_batch_size=16 +``` + +## Train Llama3-8B with SimpleFSDP + torch.compile + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.experiments.graph_trainer.llama3 \ + --config simple_fsdp_llama3_8b \ + --dataloader.dataset=c4_test \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --parallelism.context_parallel_degree=1 \ + --training.global_batch_size=16 +``` + +## Multi-node training + +```shell +export MASTER_ADDR= +export CCOM_SOCKET_IFNAME=eth0 +``` + +### Llama3 8B with HSDP + +Runs on four nodes. With torch.compile: + +```bash +torchrun --nnodes=4 --node_rank=${NODE_RANK} \ + --master_addr=${MASTER_ADDR} --master_port=29500 \ + --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.models.llama3 \ + --config llama3_8b_fullac_compile \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=16 \ + --parallelism.data_parallel_replicate_degree=4 \ + --parallelism.context_parallel_degree=1 \ + --training.global_batch_size=128 +``` + +### Llama3 70B + +Download the tokenizer: + +```shell +python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-70B --assets tokenizer --hf_token=... +``` + +Eager: + +```bash +torchrun --nnodes=4 --node_rank=${NODE_RANK} \ + --master_addr=${MASTER_ADDR} --master_port=29500 \ + --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.models.llama3 \ + --config llama3_70b_fullac_eager \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=64 \ + --parallelism.context_parallel_degree=1 \ + --training.local_batch_size=1 \ + --training.global_batch_size=128 \ + --training.steps=10 +``` + +torch.compile: + +```bash +torchrun --nnodes=4 --node_rank=${NODE_RANK} \ + --master_addr=${MASTER_ADDR} --master_port=29500 \ + --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.models.llama3 \ + --config llama3_70b_fullac_compile \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=64 \ + --parallelism.context_parallel_degree=1 \ + --training.local_batch_size=1 \ + --training.global_batch_size=128 \ + --training.steps=10 +``` + +SimpleFSDP + torch.compile: + +```bash +torchrun --nnodes=4 --node_rank=${NODE_RANK} \ + --master_addr=${MASTER_ADDR} --master_port=29500 \ + --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.experiments.graph_trainer.llama3 \ + --config simple_fsdp_llama3_70b_fullac \ + --dataloader.dataset=c4_test \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=64 \ + --parallelism.context_parallel_degree=1 \ + --training.local_batch_size=1 \ + --training.global_batch_size=128 \ + --training.steps=10 +``` + +## Host collectives (HostCC) + +To run any config above with host collectives, export these before the command: + +```shell +export TORCH_NEURONX_ENABLE_HOST_CC=1 +export TORCH_NEURONX_ENABLE_STREAMS=1 +``` + +On more than one node, also set the UltraServer mode. Use `1` for a multiple of +four nodes and `3` for two nodes only: + +```shell +export NEURON_RT_ULTRASERVER_MODE=1 +``` + +See the [UltraServer +documentation](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/deploy/eks/ultraserver.html) +for more detail. + +The first training step will take a few minutes, but after that, expect +consistent throughput. diff --git a/torchtitan_neuron/docs/moe-tutorials.md b/torchtitan_neuron/docs/moe-tutorials.md new file mode 100644 index 00000000000..9cd562b2f50 --- /dev/null +++ b/torchtitan_neuron/docs/moe-tutorials.md @@ -0,0 +1,58 @@ +# Training MoE models with torchtitan-neuron + +## Setup + +```shell +cd /workspace/torchtitan_neuron +pip install -r requirements.txt +pip install -e . +``` + +## Train Qwen3-30B-A3B with SimpleFSDP + torch.compile + +Download the tokenizer: + +```shell +python scripts/download_hf_assets.py --repo_id Qwen/Qwen3-30B-A3B --assets tokenizer +``` + +Run: + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.experiments.graph_trainer.qwen3 \ + --config qwen3_30b_a3b_nki_fullac_simplefsdp \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --parallelism.enable_async_tensor_parallel \ + --training.local_batch_size=4 \ + --training.global_batch_size=64 +``` + +## Train DeepSeek-V3-16B with SimpleFSDP + torch.compile + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.experiments.graph_trainer.deepseek_v3 \ + --config simple_fsdp_deepseek_v3_16b_nki_tp4_fsdp16 \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --parallelism.expert_parallel_degree=1 \ + --training.local_batch_size=4 \ + --training.global_batch_size=64 \ + --training.steps=20 +``` + +## Train GPT-OSS-20B with SimpleFSDP + torch.compile + +```bash +torchrun --nproc_per_node=64 --local-ranks-filter 0 -m torchtitan.train \ + --module torchtitan_neuron.experiments.graph_trainer.gpt_oss \ + --config simple_fsdp_gpt_oss_20b_nki_tp4_fsdp16 \ + --training.local_batch_size=4 \ + --training.seq_len=4096 \ + --training.steps=20 +``` + +The first training step will take a few minutes, but after that, expect +consistent throughput. diff --git a/torchtitan_neuron/docs/neuron_attention_backends.md b/torchtitan_neuron/docs/neuron_attention_backends.md new file mode 100644 index 00000000000..bee2933fba9 --- /dev/null +++ b/torchtitan_neuron/docs/neuron_attention_backends.md @@ -0,0 +1,129 @@ +# Neuron attention backends + +TorchTitan's Neuron recipes select the attention implementation for training on +AWS Trainium. Select a recipe through `--module` and `--config`; the recipe +configures attention together with the model's Neuron components and compilation +settings. + +This guide describes the attention behavior of the recipes in +`torchtitan_neuron`. Run them in an environment with PyTorch and `torch_neuronx` +configured for Neuron. Recipes using NKI kernels also require `nkilib`. + +## Backends used by Neuron recipes + +| Model | Attention backend | Behavior | +| --- | --- | --- | +| Llama3 | `sdpa` | Causal scaled dot-product attention, including grouped-query attention | +| DeepSeek-V3 | `sdpa` | Causal scaled dot-product attention | +| Qwen3 | `sdpa` | Causal scaled dot-product attention, including grouped-query attention | +| GPT-OSS | `nki` | Causal attention with native attention sinks, grouped-query attention, and per-layer sliding windows | + +The attention backend is separate from `compile.backend`. For example, a +compiled Llama3 recipe uses `sdpa` attention and the `neuron` compilation +backend. Selecting an NKI MoE recipe for DeepSeek-V3 or Qwen3 does not change +its attention backend to `nki`. + +### SDPA + +Llama3, DeepSeek-V3, and Qwen3 recipes use TorchTitan's +`ScaledDotProductAttention` wrapper around PyTorch's +`scaled_dot_product_attention`. These recipes use causal attention: a token +can attend to itself and earlier tokens in its sequence. + +The TorchTitan wrapper controls masking through `is_causal` and rejects +explicit `attention_masks`. It does not implement sliding-window attention +or packed-document isolation. This is a constraint of the TorchTitan wrapper, +not of the full PyTorch SDPA API. + +### GPT-OSS NKI attention + +The GPT-OSS recipe selects +[`NKIFlashAttentionWrapper`](../models/gpt_oss/nki_attention.py). +The wrapper uses NKI forward and backward kernels and computes gradients for +queries, keys, values, and learned attention sinks. It is integrated with +`torch.compile` through Neuron custom operators. + +GPT-OSS alternates between 128-token causal windows and full causal attention. +The wrapper accepts `sliding_window=128` or `sliding_window=None`; other window +sizes raise `ValueError` when the module is constructed. Keep the window +selection supplied by the model recipe. + +The `nki` attention backend is specific to GPT-OSS. Replacing it with `sdpa` +is not supported for its sliding-window layers: the model registry rejects +that combination. + +## Packed-document behavior + +**Neither SDPA nor GPT-OSS NKI attention enforces document boundaries within +a packed sequence.** If a sequence contains document A followed by document B, +tokens in B can attend to earlier tokens in A. A sliding window limits how far +back a token can attend, but does not isolate documents. + +TorchTitan's text dataloader packs documents into sequences. Position resets +and end-of-document tokens do not add an attention mask to these backends. +Constructing an SDPA configuration logs a warning about this behavior. + +If your training objective requires strict document isolation, use a data +pipeline that keeps each sequence within one document, or an attention +implementation that consumes document boundaries. The Neuron recipes listed +here do not provide document-boundary masking. + +## Selecting a Trainer recipe + +Pass the registry module as `--module` and the recipe name as `--config`. +The following recipes select the backends described above: + +| Registry module | Recipe | +| --- | --- | +| [`torchtitan_neuron.models.llama3.config_registry`](../models/llama3/config_registry.py) | `llama3_8b_sac_compile` | +| [`torchtitan_neuron.models.deepseek_v3.config_registry`](../models/deepseek_v3/config_registry.py) | `deepseek_v3_16b_neuron_tp4_fsdp16` | +| [`torchtitan_neuron.models.qwen3.config_registry`](../models/qwen3/config_registry.py) | `qwen3_30b_a3b_nki_fullac_compile` | +| [`torchtitan_neuron.models.gpt_oss.config_registry`](../models/gpt_oss/config_registry.py) | `gpt_oss_20b_nki_tp4_fsdp16` | + +For example, launch the GPT-OSS recipe on a single host with `torchrun`. +Set `NEURON_PROCESSES_PER_NODE` to the worker count configured for your Neuron +host and `HF_ASSETS_PATH` to the local GPT-OSS tokenizer and model assets: + +```bash +torchrun --standalone --nproc_per_node="${NEURON_PROCESSES_PER_NODE}" \ + -m torchtitan.train \ + --module torchtitan_neuron.models.gpt_oss.config_registry \ + --config gpt_oss_20b_nki_tp4_fsdp16 \ + --hf_assets_path="${HF_ASSETS_PATH}" +``` + +This recipe uses tensor parallelism of degree 4 and infers its data-parallel +shard degree from the world size. The total worker count must be divisible by +4; TP4 with FSDP16 uses 64 workers. Choose a topology with enough memory for +the model and training state. For multiple hosts, use your distributed +launcher's node and rendezvous settings with the same module and recipe. + +### GraphTrainer recipes + +The Neuron GraphTrainer recipes retain the corresponding Trainer recipe's +attention selection and use SimpleFSDP with Neuron JIT compilation: + +| Registry module | Recipe | +| --- | --- | +| [`torchtitan_neuron.experiments.graph_trainer.llama3.config_registry`](../experiments/graph_trainer/llama3/config_registry.py) | `simple_fsdp_llama3_8b` | +| [`torchtitan_neuron.experiments.graph_trainer.deepseek_v3.config_registry`](../experiments/graph_trainer/deepseek_v3/config_registry.py) | `simple_fsdp_deepseek_v3_16b_nki_tp4_fsdp16` | +| [`torchtitan_neuron.experiments.graph_trainer.qwen3.config_registry`](../experiments/graph_trainer/qwen3/config_registry.py) | `qwen3_30b_a3b_nki_fullac_simplefsdp` | +| [`torchtitan_neuron.experiments.graph_trainer.gpt_oss.config_registry`](../experiments/graph_trainer/gpt_oss/config_registry.py) | `simple_fsdp_gpt_oss_20b_nki_tp4_fsdp16` | + +Use the same training entrypoint, replacing `--module` and `--config` with +the GraphTrainer selection. + +## Configuring a training workload + +Dataset and schedule defaults vary by recipe. The Trainer recipes listed above +use the bundled `c4_test` dataset; Llama3 GraphTrainer inherits `c4` from the core +Llama3 configuration. GPT-OSS and Qwen3 NKI recipes default to 10 training steps; +DeepSeek-V3 defaults to 20. Set +`--dataloader.dataset`, `--hf_assets_path`, `--training.steps`, batch sizes, +sequence length, learning-rate schedule, checkpoint settings, and +`--dump_folder` for your workload. The recipe defaults do not define a full +training schedule. + +See [dataset configuration](../../docs/datasets.md) for input data, +[checkpointing](../../docs/checkpoint.md) for saving and resuming training, +and [Qwen3 NKI MoE](qwen3_nki_moe.md) for Qwen3 recipe details. diff --git a/torchtitan_neuron/docs/neuron_mxfp8.md b/torchtitan_neuron/docs/neuron_mxfp8.md new file mode 100644 index 00000000000..b15287df202 --- /dev/null +++ b/torchtitan_neuron/docs/neuron_mxfp8.md @@ -0,0 +1,36 @@ +# Neuron MXFP8 (trn3) + +MXFP8 dense-Linear training for qwen3 on AWS Trainium trn3 (NeuronCore-v4), on +torchao with Neuron MXFP8 support. titan builds `MXFP8Linear` with +`KernelPreference.NEURON`; torchao selects the GEMM backend internally. + +## Requirements +Trn3 hardware, torchao with Neuron MXFP8 support + +## Usage + +Torchtitan swaps each in-layer `Linear` module (attention + FFN; the LM head stays +bf16) for torchao's `MXFP8Linear` with `kernel_preference="neuron"`. MXFP8 is gated +on trn3 (`has_neuron_capability("mxfp8")`); off trn3 the converter/linear raises a +clear `ValueError`. Enable it with the named config: + +### Named config + +`torchtitan_neuron/models/qwen3/config_registry.py` ships `qwen3_8b_mxfp8`, which +attaches `MXFP8LinearConverter(fqns=["layers"], kernel_preference="neuron")` to the +model spec: + +```bash +NGPU=64 MODULE=torchtitan_neuron.models.qwen3 CONFIG=qwen3_8b_mxfp8 ./run_train.sh +``` + +## Tests + +trn3-gated unit tests: + +```bash +python -m pytest \ + torchtitan_neuron/tests/test_mxfp8_numerics.py \ + torchtitan_neuron/tests/test_capability.py \ + torchtitan_neuron/tests/test_metrics_mfu.py +``` diff --git a/torchtitan_neuron/docs/qwen3_235b_full_ep128.md b/torchtitan_neuron/docs/qwen3_235b_full_ep128.md new file mode 100644 index 00000000000..7c1dc95488d --- /dev/null +++ b/torchtitan_neuron/docs/qwen3_235b_full_ep128.md @@ -0,0 +1,147 @@ +# Qwen3-235B training on a four-node Trainium3 UltraServer + +This recipe runs 50 training steps of the full 94-layer Qwen3-235B-A22B model using +TorchTitan, FSDP2, and full activation checkpointing (FullAC). + +It uses TorchTitan's bundled tokenizer and C4 sample, starting from randomly +initialized model weights. + +## Prerequisites + +Start with a working four-node Trainium3 UltraServer using your image on all four +nodes. Use the image's `python` and `torchrun` from `PATH`. + +Open a Bash shell inside each node's container. Paste the command blocks below +into those shells in order, keeping the same shell open so the exported variables +carry into the launch command. You do not need to create a Bash script. +No package installation, branch checkout, or cherry-pick is needed. + +| Setting    | Value | +| --- | --- | +| Model    | Qwen3-235B-A22B, all 94 layers | +| Parallelism    | EP128, TP4, FSDP64 | +| Sequence length    | 4,096 tokens | +| Batch size    | 1 local, 64 global | +| Activation checkpointing    | FullAC | +| Training steps    | 50 | + +The four nodes run 256 processes. Dense parameters use FSDP64 and TP4. For expert +parameters, EP128 distributes the 128 experts across 128 positions, and FSDP +splits each expert's weights across two ranks. FSDP gathers those weights when +needed for computation. This is the named configuration's standard layout. + +## 1. Select the tokenizer and dataset + +The tokenizer and C4 sample are included in your image. Set these paths on every +node before proceeding to the launch environment. + +```bash +mkdir -p /workspace/outputs +export ASSET_ROOT=/workspace/torchtitan/tests/assets +export TOKENIZER_PATH="$ASSET_ROOT/tokenizer" +``` + +The loader tokenizes and packs the C4 sample into training sequences, repeating +the sample as needed. + +## 2. Resolve node topology + +Run this inside each node's container to identify its physical UltraServer position: + +```bash +neuron-ls --topology +``` + +Confirm all four outputs have the same `ultraserver-id`, `ultraserver-size: 4`, +and distinct `ultraserver-instance-index` values covering `0`, `1`, `2`, and `3`. +On each node, set `NODE_RANK` to its reported index: + +```bash +export NODE_RANK=0 # Use this node's ultraserver-instance-index: 0, 1, 2, or 3. +``` + +For this four-node Trn3 layout, physical indices `0` and `1` share one rack, +and `2` and `3` share the other. With 64 processes per node, the configuration +forms EP128 groups from node ranks `0` + `1` and `2` + `3`, keeping each EP +group within a rack. FSDP communication can still cross racks. Use the physical +indices, not hostname or container ordering; `torchrun` does not choose this +placement for you. + +## 3. Set the launch environment + +On every node, use the same reachable address of the node whose +`ultraserver-instance-index` is `0`, and the same master port. +Keep the `NODE_RANK` assigned in the previous step. + +```bash +export MASTER_ADDR='NODE_0_REACHABLE_ADDRESS' +export MASTER_PORT=29500 + +export TORCH_NEURONX_FALLBACK_ONLY_FOR_UNIMPLEMENTED_OPS=0 +export NEURON_CC_FLAGS="--enable-mixed-precision-accumulation --enable-saturate-infinity --auto-cast=none" +export CCOM_SOCKET_IFNAME=eth0 +export NEURON_RT_DBG_MAX_A2AV_BLOCK_SIZE=8388608 +``` + +Set `CCOM_SOCKET_IFNAME` to the reachable network interface if it differs from `eth0`. + +**Why set `NEURON_RT_DBG_MAX_A2AV_BLOCK_SIZE=8388608`?** This performance knob +specifies an upper bound on the A2A-v payload per peer, in bytes. For this +configuration, 8 MiB avoids scheduling unnecessary transfer iterations. + +Set the flag on every rank. Recalculate the bound if you change the batch, sequence +length, precision, or parallelism: an undersized value can omit required transfers. + +## 4. Launch once on each node + +Paste this command into each of the four prepared shells, using the respective +`NODE_RANK` set above. Start all four launches; the first waits for the others to join. + +`-m torchtitan.train` selects the installed Python entrypoint. +`--module` selects TorchTitan's configuration registry. + +```bash +cd /workspace/torchtitan || exit 1 +torchrun \ + --nnodes=4 \ + --nproc-per-node=64 \ + --node-rank="$NODE_RANK" \ + --master-addr="$MASTER_ADDR" \ + --master-port="$MASTER_PORT" \ + --max-restarts=0 \ + -m torchtitan.train \ + --module torchtitan_neuron.models.qwen3.config_registry \ + --config qwen3_235b_94layers_ep128_tp4_fsdp64_pp1_4k_neuron_ep_compile \ + --training.local_batch_size=1 \ + --training.global_batch_size=64 \ + --training.seq_len=4096 \ + --training.dtype=float32 \ + --training.mixed_precision_param=bfloat16 \ + --parallelism.enable_sequence_parallel \ + --debug.seed=42 \ + --training.steps=50 \ + --lr_scheduler.warmup_steps=2 \ + --metrics.log_freq=1 \ + --dataloader.dataset=c4_test \ + --dataloader.dataset_path="$ASSET_ROOT/c4_test" \ + --hf_assets_path="$TOKENIZER_PATH" \ + --dump_folder=/workspace/outputs/qwen235b-full94 +``` + +The named configuration supplies FullAC, the parallelism degrees, and Neuron +model/loss compilation. The command sets the run length to 50 steps and the seed +to 42. Checkpoint saving is disabled by default; ordinary Titan logging remains +enabled. + +## 5. Read the result + +Check the startup configuration for 256 ranks, 94 layers, global batch 64, +sequence length 4,096, and FullAC. + +Initial compilation can make early steps substantially slower. Use Titan's +standard loss, gradient-norm, throughput, step-time, and MFU records, and identify +the steady-state step window when reporting performance. Learning-rate warmup +alone does not establish that compilation has finished. + +Record the delivered image digest, tokenizer/dataset paths, and any configuration +changes alongside results. diff --git a/torchtitan_neuron/docs/qwen3_inference.md b/torchtitan_neuron/docs/qwen3_inference.md new file mode 100644 index 00000000000..4b196b1c35e --- /dev/null +++ b/torchtitan_neuron/docs/qwen3_inference.md @@ -0,0 +1,69 @@ +# Qwen3 Inference with torch.compile using TorchTitan + +## Setup TorchTitan +```shell +cd /workspace/torchtitan_neuron +pip install -r requirements.txt +pip install -e . +pip install transformers +``` + +## Download Qwen3 Model Assets +```shell +# Download model weights and tokenizer +python scripts/download_hf_assets.py --repo_id Qwen/Qwen3-8B --assets tokenizer safetensors config +``` + +## Inference with torch.compile + +Create the inference script `qwen3_inference.py`: + +```python +import torch +import torch_neuronx +from transformers import AutoModelForCausalLM, AutoTokenizer +from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel +from torch.distributed.device_mesh import init_device_mesh + +# Load Qwen3 model and tokenizer +model_name = "Qwen/Qwen3-8B" +model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16) +tokenizer = AutoTokenizer.from_pretrained(model_name) + +# Apply tensor parallel based on NEURON_RT_NUM_CORES +import os +torch.distributed.init_process_group() +num_cores = torch.distributed.get_world_size() +tp_mesh = init_device_mesh("neuron", (num_cores,)) +assert tp_mesh is not None, "Mesh initialization failed" +parallel_plan = {} +for name, module in model.named_modules(): + if "mlp.gate_proj" in name or "mlp.up_proj" in name: + parallel_plan[name] = ColwiseParallel() + elif "mlp.down_proj" in name: + parallel_plan[name] = RowwiseParallel() +model = parallelize_module(model, tp_mesh, parallel_plan) +model = model.to('neuron') + +# Compile model with neuron backend +compiled_model = torch.compile(model, backend="neuron") + +# Prepare input +text = "Hello, how are you?" +input_ids = tokenizer.encode(text, return_tensors="pt").to('neuron') + +# Run inference +with torch.no_grad(): + output = compiled_model.generate(input_ids, max_length=50, do_sample=False) + response = tokenizer.decode(output[0], skip_special_tokens=True) + print(f"Response: {response}") +``` + +## Environment Setup and Inference +```shell +export NEURON_RT_NUM_CORES=32 # Adjust based on your instance +export NEURON_CC_FLAGS="--model-type=transformer" # Trigger compiler heuristics better tuned for Transformer model architecture +torchrun --nproc-per-node=$NEURON_RT_NUM_CORES qwen3_inference.py +``` + +The first inference will trigger compilation which may take longer. Subsequent runs will use cached graphs for faster execution. diff --git a/torchtitan_neuron/docs/qwen3_nki_moe.md b/torchtitan_neuron/docs/qwen3_nki_moe.md new file mode 100644 index 00000000000..77161aa95a6 --- /dev/null +++ b/torchtitan_neuron/docs/qwen3_nki_moe.md @@ -0,0 +1,77 @@ +# Qwen3 NKI MoE + +The Neuron Qwen3 registries replace the standard mixture-of-experts +implementation with the Qwen3 NKI MoE kernel. They provide recipes for both +Trainer and GraphTrainer. + +## Neuron components + +`torchtitan_neuron.models.qwen3.config_registry` constructs Qwen3 with: + +- [causal SDPA attention](neuron_attention_backends.md#sdpa) +- the sequential RoPE prefix-cache path, which avoids a per-position gather +- `NKIMoE.Config` in place of the standard MoE implementation +- the compiled gradient-norm override when model compilation is enabled + +The launchable NKI MoE recipes documented below use the Qwen3 `30B-A3B` +model preset. + +Attention remains SDPA when NKI MoE is enabled. Documents packed into the +same sequence can attend across document boundaries; see +[packed-document behavior](neuron_attention_backends.md#packed-document-behavior) +before choosing a data pipeline. + +## Trainer recipes + +Run Trainer with model compilation disabled: + +```bash +python -m torchtitan.train \ + --module torchtitan_neuron.models.qwen3.config_registry \ + --config qwen3_30b_a3b_nki_fullac_eager \ + --hf_assets_path=./assets/hf/Qwen3-30B-A3B +``` + +Run Trainer with the Neuron compilation backend: + +```bash +python -m torchtitan.train \ + --module torchtitan_neuron.models.qwen3.config_registry \ + --config qwen3_30b_a3b_nki_fullac_compile \ + --hf_assets_path=./assets/hf/Qwen3-30B-A3B +``` + +## GraphTrainer recipe + +The GraphTrainer recipe uses SimpleFSDP and Neuron JIT compilation: + +```bash +python -m torchtitan.train \ + --module torchtitan_neuron.experiments.graph_trainer.qwen3.config_registry \ + --config qwen3_30b_a3b_nki_fullac_simplefsdp \ + --hf_assets_path=./assets/hf/Qwen3-30B-A3B \ + --parallelism.tensor_parallel_degree=4 \ + --parallelism.data_parallel_shard_degree=-1 \ + --parallelism.enable_async_tensor_parallel +``` + +## Global-batch load-balancing loss + +`qwen3_30b_a3b_nki_global_lbl_fullac_simplefsdp` trains with Qwen3's +global-batch load-balancing loss (`load_balancing="global_aux_loss"`) instead +of the default expert bias. Each MoE layer adds +`load_balance_coeff * E * sum_i f_i * P_i`, where `f_i` is the fraction of +routed tokens assigned to expert `i` (from token counts summed over the data- +and context-parallel ranks) and `P_i` is expert `i`'s mean router probability +over the global batch. To enable it on another +Qwen3 configuration, apply +`torchtitan.models.qwen3.config_registry.use_global_batch_lbl`. It is refused +with a non-default `spmd_backend`, GraphTrainer EP overlap, and precompile. + +## Training configuration + +These recipes use the bundled `c4_test` dataset, 10 training steps, 10 warmup +steps, and per-step metrics so that a new installation can perform a bounded +smoke run. Before starting a training workload, override the dataset, asset +path, step count, checkpoint, output, and topology settings for the target +environment. diff --git a/torchtitan_neuron/experiments/__init__.py b/torchtitan_neuron/experiments/__init__.py new file mode 100644 index 00000000000..f1f8411ed15 --- /dev/null +++ b/torchtitan_neuron/experiments/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron-specific TorchTitan experiment registries.""" diff --git a/torchtitan_neuron/experiments/graph_trainer/__init__.py b/torchtitan_neuron/experiments/graph_trainer/__init__.py new file mode 100644 index 00000000000..64a9633eaf6 --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron-specific GraphTrainer registries.""" diff --git a/torchtitan_neuron/experiments/graph_trainer/configs.py b/torchtitan_neuron/experiments/graph_trainer/configs.py new file mode 100644 index 00000000000..ee61653c22e --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/configs.py @@ -0,0 +1,43 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Shared configuration helpers for Neuron GraphTrainer recipes.""" + +from collections.abc import Callable +from dataclasses import replace + +from torchtitan.experiments.graph_trainer.configs import ( + GraphTrainerCompileConfig, + to_graph_trainer_config, +) +from torchtitan.experiments.graph_trainer.trainer import GraphTrainer +from torchtitan.protocols.model_spec import ModelSpec +from torchtitan.trainer import Trainer + + +def neuron_simple_fsdp( + base_config: Trainer.Config, + model_registry: Callable[[str], ModelSpec], + *, + model_name: str, + model_spec: ModelSpec | None = None, +) -> GraphTrainer.Config: + """Convert a Neuron Trainer recipe to a JIT SimpleFSDP recipe.""" + activation_checkpoint = base_config.activation_checkpoint + compile_config = base_config.compile + graph_config = to_graph_trainer_config(base_config, model_registry) + graph_config.model_spec = ( + model_spec + if model_spec is not None + else replace(graph_config.model_spec, name=model_name) + ) + graph_config.activation_checkpoint = activation_checkpoint + graph_config.compile = GraphTrainerCompileConfig( + enable=compile_config.enable, + components=list(compile_config.components), + backend=compile_config.backend, + mode="jit", + ) + return graph_config diff --git a/torchtitan_neuron/experiments/graph_trainer/deepseek_v3/__init__.py b/torchtitan_neuron/experiments/graph_trainer/deepseek_v3/__init__.py new file mode 100644 index 00000000000..df8458c0509 --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/deepseek_v3/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""DeepSeek-V3 Neuron GraphTrainer recipes.""" diff --git a/torchtitan_neuron/experiments/graph_trainer/deepseek_v3/config_registry.py b/torchtitan_neuron/experiments/graph_trainer/deepseek_v3/config_registry.py new file mode 100644 index 00000000000..895e97813c5 --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/deepseek_v3/config_registry.py @@ -0,0 +1,34 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""DeepSeek-V3 recipes that require NKI kernels and the Neuron JIT path.""" + +from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC +from torchtitan.experiments.graph_trainer.deepseek_v3 import model_registry +from torchtitan.experiments.graph_trainer.trainer import GraphTrainer + +from torchtitan_neuron.experiments.graph_trainer.configs import neuron_simple_fsdp +from torchtitan_neuron.models.deepseek_v3.config_registry import ( + deepseek_v3_16b_neuron_tp4_fsdp16, +) + + +def simple_fsdp_deepseek_v3_16b_nki_tp4_fsdp16() -> GraphTrainer.Config: + config = neuron_simple_fsdp( + deepseek_v3_16b_neuron_tp4_fsdp16(), + model_registry, + model_name="simple_fsdp/deepseek_v3", + ) + config.training.local_batch_size = 4 + config.training.global_batch_size = 64 + config.training.seq_len = 4096 + config.activation_checkpoint = SelectiveAC.Config() + return config + + +def simple_fsdp_deepseek_v3_16b_nki_tp4_fsdp16_fac() -> GraphTrainer.Config: + config = simple_fsdp_deepseek_v3_16b_nki_tp4_fsdp16() + config.activation_checkpoint = FullAC.Config() + return config diff --git a/torchtitan_neuron/experiments/graph_trainer/gpt_oss/__init__.py b/torchtitan_neuron/experiments/graph_trainer/gpt_oss/__init__.py new file mode 100644 index 00000000000..fd2d7247c1c --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/gpt_oss/__init__.py @@ -0,0 +1,38 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +from dataclasses import fields + +from torchtitan.components.optimizer import register_moe_load_balancing_hook +from torchtitan.experiments.graph_trainer.graph_pp.pipeline import graph_pipeline_llm +from torchtitan.models.gpt_oss import gptoss_configs +from torchtitan.models.gpt_oss.state_dict_adapter import GptOssStateDictAdapter +from torchtitan.protocols.model_spec import ModelSpec + +from .model import GraphTrainerGptOssModel +from .parallelize import parallelize_gptoss + + +def model_registry( + flavor: str, + moe_comm_backend: str = "standard", + attn_backend: str = "varlen", +) -> ModelSpec: + base = gptoss_configs[flavor]( + moe_comm_backend=moe_comm_backend, + attn_backend=attn_backend, + ) + config = GraphTrainerGptOssModel.Config( + **{f.name: getattr(base, f.name) for f in fields(base)} + ) + return ModelSpec( + name="graph_trainer/gpt_oss", + flavor=flavor, + model=config, + parallelize_fn=parallelize_gptoss, + pipelining_fn=graph_pipeline_llm, + post_optimizer_build_fn=register_moe_load_balancing_hook, + state_dict_adapter=GptOssStateDictAdapter, + ) diff --git a/torchtitan_neuron/experiments/graph_trainer/gpt_oss/config_registry.py b/torchtitan_neuron/experiments/graph_trainer/gpt_oss/config_registry.py new file mode 100644 index 00000000000..b77956335a2 --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/gpt_oss/config_registry.py @@ -0,0 +1,26 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""GPT-OSS recipes that require NKI kernels and the Neuron JIT path.""" + +from torchtitan.distributed.activation_checkpoint import FullAC +from torchtitan.experiments.graph_trainer.trainer import GraphTrainer + +from torchtitan_neuron.experiments.graph_trainer.configs import neuron_simple_fsdp +from torchtitan_neuron.models.gpt_oss.config_registry import gpt_oss_20b_nki_tp4_fsdp16 + +from . import model_registry + + +def simple_fsdp_gpt_oss_20b_nki_tp4_fsdp16() -> GraphTrainer.Config: + config = neuron_simple_fsdp( + gpt_oss_20b_nki_tp4_fsdp16(), + model_registry, + model_name="simple_fsdp/gpt_oss", + ) + config.training.local_batch_size = 4 + config.training.seq_len = 4096 + config.activation_checkpoint = FullAC.Config() + return config diff --git a/torchtitan_neuron/experiments/graph_trainer/gpt_oss/model.py b/torchtitan_neuron/experiments/graph_trainer/gpt_oss/model.py new file mode 100644 index 00000000000..a929a32da0d --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/gpt_oss/model.py @@ -0,0 +1,24 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +from dataclasses import dataclass + +from torchtitan.experiments.graph_trainer.simple_fsdp import ( + disable_active_parametrization, +) +from torchtitan.models.gpt_oss import GptOssModel + + +class GraphTrainerGptOssModel(GptOssModel): + @dataclass(kw_only=True, slots=True) + class Config(GptOssModel.Config): + pass + + def __init__(self, config: Config): + super().__init__(config) + + def init_weights(self, *args, **kwargs): + with disable_active_parametrization(): + super().init_weights(*args, **kwargs) diff --git a/torchtitan_neuron/experiments/graph_trainer/gpt_oss/parallelize.py b/torchtitan_neuron/experiments/graph_trainer/gpt_oss/parallelize.py new file mode 100644 index 00000000000..e2b48f327c2 --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/gpt_oss/parallelize.py @@ -0,0 +1,77 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +from torchtitan.config import ParallelismConfig, TrainingConfig +from torchtitan.distributed import ParallelDims +from torchtitan.distributed.activation_checkpoint import ActivationCheckpointingConfig +from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp +from torchtitan.experiments.graph_trainer.common_utils import ( + annotate_module_fqns, + annotate_moe_ep_regions, + apply_simple_fsdp, +) +from torchtitan.experiments.graph_trainer.compile import apply_compile +from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig + +from .model import GraphTrainerGptOssModel + + +def annotate_gptoss(model: GraphTrainerGptOssModel) -> None: + """Attach module FQN annotations for graph-trainer passes.""" + if any(layer.moe is not None for layer in model.config.layers): + annotate_moe_ep_regions() + annotate_module_fqns(model) + + +def parallelize_gptoss( + model: GraphTrainerGptOssModel, + *, + parallel_dims: ParallelDims, + training: TrainingConfig, + parallelism: ParallelismConfig, + compile_config: GraphTrainerCompileConfig, + ac_config: ActivationCheckpointingConfig, + dump_folder: str, +): + assert ( + not parallel_dims.cp_enabled + ), "GraphTrainer GPT-OSS NKI does not currently support context parallelism" + assert ( + not parallel_dims.ep_enabled + ), "GraphTrainer GPT-OSS NKI does not currently support expert parallelism" + + assert ( + training.seq_len % parallel_dims.seq_len_divisor == 0 + ), f""" + Sequence length {training.seq_len} must be divisible by the product of TP degree + ({parallel_dims.tp}) and 2 * CP degree ({parallel_dims.cp}), i.e. {parallel_dims.seq_len_divisor}. + """ + + annotate_gptoss(model) + + if parallel_dims.tp_enabled: + model.parallelize(parallel_dims) + + if parallel_dims.tp_enabled: + maybe_enable_async_tp(parallelism, compile_config, parallel_dims.get_mesh("tp")) + + if compile_config.mode == "jit" and ac_config is not None: + # JIT mode does not run graph-trainer's AOT memory-policy passes, so + # apply the eager activation-checkpoint policy before SimpleFSDP. + ac_config.build(dump_folder=dump_folder).apply(model) + + # Apply SimpleFSDP unconditionally. The fsdp mesh exists even at degree 1, + # which lets MixedPrecisionPolicy still cast parameters in single-rank runs. + model = apply_simple_fsdp(model, parallel_dims=parallel_dims, training=training) + + model = apply_compile( + model, + compile_config=compile_config, + parallelism=parallelism, + parallel_dims=parallel_dims, + dump_folder=dump_folder, + ) + + return model diff --git a/torchtitan_neuron/experiments/graph_trainer/llama3/__init__.py b/torchtitan_neuron/experiments/graph_trainer/llama3/__init__.py new file mode 100644 index 00000000000..1faf16ab3bb --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/llama3/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Llama3 Neuron GraphTrainer configs.""" diff --git a/torchtitan_neuron/experiments/graph_trainer/llama3/config_registry.py b/torchtitan_neuron/experiments/graph_trainer/llama3/config_registry.py new file mode 100644 index 00000000000..aefb548ba93 --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/llama3/config_registry.py @@ -0,0 +1,70 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Llama3 SimpleFSDP recipes for the Neuron JIT GraphTrainer path.""" + +from torchtitan.components.quantization import MXFP8LinearConverter +from torchtitan.config import CompileConfig +from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC +from torchtitan.experiments.graph_trainer.llama3 import model_registry +from torchtitan.experiments.graph_trainer.trainer import GraphTrainer +from torchtitan.models.llama3.config_registry import llama3_70b, llama3_8b +from torchtitan.protocols.model import ModelConfigConverter + +from torchtitan_neuron.experiments.graph_trainer.configs import neuron_simple_fsdp + + +def _neuron_simple_fsdp_llama3_config( + config_factory, + *, + full_ac: bool = False, + converters: list[ModelConfigConverter.Config] | None = None, +) -> GraphTrainer.Config: + base = config_factory() + base.model_spec = model_registry( + base.model_spec.flavor, attn_backend="sdpa", converters=converters + ) + # Neuron does not support the RoPE gather; use the prefix-cache positions path. + for layer in base.model_spec.model.layers: + layer.attention.rope.use_prefix_cache_for_positions = True + base.activation_checkpoint = FullAC.Config() if full_ac else SelectiveAC.Config() + base.compile = CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ) + config = neuron_simple_fsdp(base, model_registry, model_name="simple_fsdp/llama3") + return config + + +def simple_fsdp_llama3_8b() -> GraphTrainer.Config: + return _neuron_simple_fsdp_llama3_config(llama3_8b) + + +def simple_fsdp_llama3_8b_fullac() -> GraphTrainer.Config: + return _neuron_simple_fsdp_llama3_config(llama3_8b, full_ac=True) + + +def simple_fsdp_llama3_70b_fullac() -> GraphTrainer.Config: + return _neuron_simple_fsdp_llama3_config(llama3_70b, full_ac=True) + + +def simple_fsdp_llama3_8b_mxfp8() -> GraphTrainer.Config: + """Llama3-8B SimpleFSDP TP4 x FSDP16 with MXFP8 linears. + + MXFP8 counterpart of ``simple_fsdp_llama3_8b``, same ``8B`` flavor. + ``fqns=["layers"]`` converts 160 linears (per layer: ``wqkv``, ``wo``, + ``w1``/``w2``/``w3``); ``lm_head`` and embeddings stay bf16. + """ + return _neuron_simple_fsdp_llama3_config( + llama3_8b, + converters=[ + MXFP8LinearConverter.Config( + fqns=["layers"], + kernel_preference="neuron", + model_compile_enabled=True, + ) + ], + ) diff --git a/torchtitan_neuron/experiments/graph_trainer/qwen3/__init__.py b/torchtitan_neuron/experiments/graph_trainer/qwen3/__init__.py new file mode 100644 index 00000000000..64a251d95c6 --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/qwen3/__init__.py @@ -0,0 +1,41 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Qwen3 Neuron GraphTrainer configs.""" + +from dataclasses import fields + +from torchtitan.experiments.graph_trainer.graph_pp.pipeline import graph_pipeline_llm +from torchtitan.experiments.graph_trainer.qwen3.model import GraphTrainerQwen3Model +from torchtitan.experiments.graph_trainer.qwen3.parallelize import parallelize_qwen3 +from torchtitan.models.qwen3.state_dict_adapter import Qwen3StateDictAdapter +from torchtitan.protocols.model_spec import ModelSpec + +from torchtitan_neuron.models.qwen3 import model_registry as qwen3_model_registry + + +def model_registry( + flavor: str, + attn_backend: str = "flex", + moe_comm_backend: str | None = None, +) -> ModelSpec: + """Build GraphTrainer specs for core and Neuron-owned Qwen3 flavors.""" + base = qwen3_model_registry( + flavor, + attn_backend=attn_backend, + moe_comm_backend=moe_comm_backend, + ).model + config = GraphTrainerQwen3Model.Config( + **{field.name: getattr(base, field.name) for field in fields(base)} + ) + return ModelSpec( + name="graph_trainer/qwen3", + flavor=flavor, + model=config, + parallelize_fn=parallelize_qwen3, + pipelining_fn=graph_pipeline_llm, + post_optimizer_build_fn=None, + state_dict_adapter=Qwen3StateDictAdapter, + ) diff --git a/torchtitan_neuron/experiments/graph_trainer/qwen3/config_registry.py b/torchtitan_neuron/experiments/graph_trainer/qwen3/config_registry.py new file mode 100644 index 00000000000..129acf27048 --- /dev/null +++ b/torchtitan_neuron/experiments/graph_trainer/qwen3/config_registry.py @@ -0,0 +1,375 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Qwen3 NKI MoE SimpleFSDP Neuron configs.""" + +from dataclasses import fields, replace + +from torchtitan.components.metrics import MetricsProcessor +from torchtitan.components.optimizer import ( + default_adamw, + register_moe_load_balancing_hook, +) +from torchtitan.components.quantization import MXFP8LinearConverter +from torchtitan.config import CompileConfig +from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC +from torchtitan.distributed.pipeline_parallel import pipeline_llm +from torchtitan.experiments.graph_trainer.configs import ( + GraphTrainerCompileConfig, + to_graph_trainer_config, +) +from torchtitan.experiments.graph_trainer.qwen3.model import GraphTrainerQwen3Model +from torchtitan.experiments.graph_trainer.qwen3.parallelize import parallelize_qwen3 +from torchtitan.experiments.graph_trainer.trainer import GraphTrainer +from torchtitan.models.qwen3 import model_registry as dense_qwen3_model_registry +from torchtitan.models.qwen3.config_registry import qwen3_32b, use_global_batch_lbl +from torchtitan.protocols.model import ModelConfigConverter +from torchtitan.protocols.model_spec import ModelSpec +from torchtitan.trainer import Trainer + +from torchtitan_neuron.components.quantization.mx_prequant import ( + NeuronMXFP8LinearConverter, +) +from torchtitan_neuron.experiments.graph_trainer.configs import neuron_simple_fsdp +from torchtitan_neuron.models.qwen3._mxfp8_flavors import register_mxfp8_flavors +from torchtitan_neuron.models.qwen3.config_registry import ( + qwen3_235b_4layers_ep64_nle1_tp4_neuron_ep, + qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_compile, + qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_pp2_compile, + qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_compile, + qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_pp2_compile, + qwen3_235b_scaled_64rank_noep, + qwen3_nki_moe_config, +) + +from . import model_registry as graph_model_registry + + +def graph_trainer_qwen3_235b_scaled() -> GraphTrainer.Config: + config = to_graph_trainer_config( + qwen3_235b_scaled_64rank_noep(), + graph_model_registry, + ) + config.compile = GraphTrainerCompileConfig(enable=True) + return config + + +def graph_trainer_qwen3_235b_4layers_ep64_nle1_tp4_neuron_ep() -> (GraphTrainer.Config): + config = to_graph_trainer_config( + qwen3_235b_4layers_ep64_nle1_tp4_neuron_ep(), + graph_model_registry, + ) + config.compile = GraphTrainerCompileConfig(enable=True) + return config + + +def _neuron_qwen3_simple_fsdp_config( + flavor: str, + *, + tensor_parallel_degree: int, + global_batch_size: int, + selective_ac: bool = False, + lr: float | None = None, + converters: list[ModelConfigConverter.Config] | None = None, +) -> GraphTrainer.Config: + base = qwen3_32b() + if lr is not None: + base.optimizer = default_adamw(lr=lr) + base.model_spec = dense_qwen3_model_registry( + flavor, + attn_backend="sdpa", + converters=converters, + ) + for layer in base.model_spec.model.layers: + layer.attention.rope.use_prefix_cache_for_positions = True + base.hf_assets_path = f"./assets/hf/Qwen3-{flavor}" + base.metrics = MetricsProcessor.Config(log_freq=1) + base.dataloader.dataset = "c4_test" + base.training.local_batch_size = 1 + base.training.steps = 10 + base.training.global_batch_size = global_batch_size + base.parallelism.tensor_parallel_degree = tensor_parallel_degree + base.parallelism.data_parallel_shard_degree = -1 + base.parallelism.context_parallel_degree = 1 + base.parallelism.pipeline_parallel_degree = 1 + config = to_graph_trainer_config(base, graph_model_registry) + config.model_spec = replace(config.model_spec, name="simple_fsdp/qwen3") + config.activation_checkpoint = ( + SelectiveAC.Config() if selective_ac else FullAC.Config() + ) + config.compile = GraphTrainerCompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + mode="jit", + ) + return config + + +def simple_fsdp_qwen3_0_6b_tp1_fsdp64_fullac() -> GraphTrainer.Config: + return _neuron_qwen3_simple_fsdp_config( + "0.6B", + tensor_parallel_degree=1, + global_batch_size=64, + lr=3e-4, + ) + + +def simple_fsdp_qwen3_1_7b_tp1_fsdp64_fullac() -> GraphTrainer.Config: + return _neuron_qwen3_simple_fsdp_config( + "1.7B", + tensor_parallel_degree=1, + global_batch_size=64, + ) + + +def simple_fsdp_qwen3_4b_tp4_fsdp16_sac() -> GraphTrainer.Config: + return _neuron_qwen3_simple_fsdp_config( + "4B", + tensor_parallel_degree=4, + global_batch_size=16, + selective_ac=True, + ) + + +def simple_fsdp_qwen3_8b_tp1_fsdp64_fullac() -> GraphTrainer.Config: + return _neuron_qwen3_simple_fsdp_config( + "8B", + tensor_parallel_degree=1, + global_batch_size=64, + ) + + +def simple_fsdp_qwen3_14b_tp4_fsdp16_fullac() -> GraphTrainer.Config: + return _neuron_qwen3_simple_fsdp_config( + "14B", + tensor_parallel_degree=4, + global_batch_size=16, + ) + + +def simple_fsdp_qwen3_32b_tp4_fsdp16_fullac() -> GraphTrainer.Config: + return _neuron_qwen3_simple_fsdp_config( + "32B", + tensor_parallel_degree=4, + global_batch_size=16, + ) + + +def simple_fsdp_qwen3_32b_tp4_fsdp16_sac() -> GraphTrainer.Config: + return _neuron_qwen3_simple_fsdp_config( + "32B", + tensor_parallel_degree=4, + global_batch_size=16, + selective_ac=True, + ) + + +def qwen3_moe_simple_fsdp_neuron_model_registry( + base_config: Trainer.Config, +) -> ModelSpec: + base = base_config.model_spec.model + model_config = GraphTrainerQwen3Model.Config( + **{field.name: getattr(base, field.name) for field in fields(base)} + ) + return ModelSpec( + name="simple_fsdp/qwen3", + flavor=base_config.model_spec.flavor, + model=model_config, + parallelize_fn=parallelize_qwen3, + pipelining_fn=None, + post_optimizer_build_fn=register_moe_load_balancing_hook, + state_dict_adapter=None, + ) + + +def _to_qwen3_neuron_graph_trainer_config( + config: Trainer.Config, +) -> GraphTrainer.Config: + """Use the core Qwen3 JIT parallelizer while preserving Neuron recipe state.""" + model_spec = qwen3_moe_simple_fsdp_neuron_model_registry(config) + return neuron_simple_fsdp( + config, + lambda _flavor: model_spec, + model_name="simple_fsdp/qwen3", + model_spec=model_spec, + ) + + +def _qwen3_moe_simple_fsdp_neuron_compile( + preset_name: str, + *, + n_layers: int | None = None, + num_experts: int | None = None, + top_k: int | None = None, +) -> GraphTrainer.Config: + config = qwen3_nki_moe_config( + preset_name, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + n_layers=n_layers, + num_experts=num_experts, + top_k=top_k, + ) + return _to_qwen3_neuron_graph_trainer_config(config) + + +def qwen3_30b_a3b_nki_fullac_simplefsdp() -> GraphTrainer.Config: + return _qwen3_moe_simple_fsdp_neuron_compile("30B-A3B") + + +def qwen3_30b_a3b_nki_global_lbl_fullac_simplefsdp() -> GraphTrainer.Config: + config = qwen3_nki_moe_config( + "30B-A3B", compile_config=CompileConfig(enable=True, backend="neuron") + ) + return _to_qwen3_neuron_graph_trainer_config(use_global_batch_lbl(config)) + + +def qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_compile_simplefsdp() -> ( + GraphTrainer.Config +): + """Six-layer SimpleFSDP neuron_ep compile configuration.""" + config = qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_compile() + return _to_qwen3_neuron_graph_trainer_config(config) + + +def qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_compile_simplefsdp() -> ( + GraphTrainer.Config +): + """EP128/4K SimpleFSDP counterpart of the FSDP2 configuration.""" + config = qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_compile() + return _to_qwen3_neuron_graph_trainer_config(config) + + +def qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_pp2_compile_simplefsdp() -> ( + GraphTrainer.Config +): + """pp=2 (1F1B) SimpleFSDP variant of the reduced-EP config. + + Inherits pp=2 / 1F1B / lbs=16 from the FSDP2 pp2 config, then swaps in the + SimpleFSDP model_spec. That fresh spec defaults pipelining_fn=None, so we + restore PP here + """ + config = qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_pp2_compile() + config = _to_qwen3_neuron_graph_trainer_config(config) + config.model_spec = replace(config.model_spec, pipelining_fn=pipeline_llm) + return config + + +def qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_pp2_compile_simplefsdp() -> ( + GraphTrainer.Config +): + """pp=2 (1F1B) SimpleFSDP variant of the EP128/4K config (4 Trn3 nodes). + + Inherits pp=2 / 1F1B / lbs=16 from the FSDP2 ep128 pp2 config, then swaps in the + SimpleFSDP model_spec. That fresh spec defaults pipelining_fn=None, so we + restore PP here + """ + config = qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_pp2_compile() + config = _to_qwen3_neuron_graph_trainer_config(config) + config.model_spec = replace(config.model_spec, pipelining_fn=pipeline_llm) + return config + + +# --- Dense Qwen3 SimpleFSDP MXFP8 (reuses the Neuron SimpleFSDP helper; +# dim6144/dim8192 flavors are registered by +# torchtitan_neuron.models.qwen3._mxfp8_flavors) --- + + +def simple_fsdp_qwen3_32b_tp4_fsdp16_fullac_mxfp8() -> GraphTrainer.Config: + """Qwen3-32B SimpleFSDP TP4 x FSDP16 FullAC, MXFP8 on all in-layer linears. + + ``fqns=["layers"]`` converts per-layer wqkv/wo/w1/w2/w3; lm_head and embeddings + stay bf16. + """ + return _neuron_qwen3_simple_fsdp_config( + "32B", + tensor_parallel_degree=4, + global_batch_size=16, + converters=[ + MXFP8LinearConverter.Config( + fqns=["layers"], + kernel_preference="neuron", + model_compile_enabled=True, + ) + ], + ) + + +def simple_fsdp_qwen3_32b_tp4_fsdp16_sac_mxfp8() -> GraphTrainer.Config: + """Qwen3-32B SimpleFSDP TP4 x FSDP16 SAC, MXFP8 on all in-layer linears. + + Selective-AC variant of the FullAC MXFP8 config; same ``fqns=["layers"]`` + conversion. + """ + return _neuron_qwen3_simple_fsdp_config( + "32B", + tensor_parallel_degree=4, + global_batch_size=16, + selective_ac=True, + converters=[ + MXFP8LinearConverter.Config( + fqns=["layers"], + kernel_preference="neuron", + model_compile_enabled=True, + ) + ], + ) + + +def simple_fsdp_qwen3_dim6144_6layers_tp4_fsdp16_fullac_mxfp8() -> GraphTrainer.Config: + """Qwen3 6-layer SimpleFSDP TP4 x FSDP16 FullAC MXFP8 on 128-aligned shapes. + + Pre-quantizes the SimpleFSDP weights with the 1x32 scheme (128-aligned per-rank + shards). Requires SimpleFSDP dim-1 sharding support. + """ + register_mxfp8_flavors() + config = _neuron_qwen3_simple_fsdp_config( + "32B_dim6144_ffn24576_6layers", + tensor_parallel_degree=4, + global_batch_size=16, + converters=[ + NeuronMXFP8LinearConverter.Config( + fqns=["layers"], + kernel_preference="neuron", + model_compile_enabled=True, + prequantize_weights=True, + prequant_scheme="1x32", + ) + ], + ) + # Random init on c4_test, no HF weights loaded, so only the tokenizer is needed and + # the in-repo one suffices (its 2016 entries fit inside the model's vocab_size). + config.hf_assets_path = "./tests/assets/tokenizer" + return config + + +def simple_fsdp_qwen3_dim8192_6layers_tp4_fsdp16_fullac_mxfp8() -> GraphTrainer.Config: + """Qwen3 6-layer SimpleFSDP TP4 x FSDP16 FullAC MXFP8 on 512-aligned shapes. + + Pre-quantizes the SimpleFSDP weights with the wrapX scheme (512-aligned per-rank + K). Requires SimpleFSDP dim-1 sharding support. + """ + register_mxfp8_flavors() + config = _neuron_qwen3_simple_fsdp_config( + "32B_dim8192_ffn32768_6layers", + tensor_parallel_degree=4, + global_batch_size=16, + converters=[ + NeuronMXFP8LinearConverter.Config( + fqns=["layers"], + kernel_preference="neuron", + model_compile_enabled=True, + prequantize_weights=True, + prequant_scheme="wrapX", + ) + ], + ) + # Random init on c4_test, so only the tokenizer is needed and the in-repo one fits. + config.hf_assets_path = "./tests/assets/tokenizer" + return config diff --git a/torchtitan_neuron/models/__init__.py b/torchtitan_neuron/models/__init__.py new file mode 100644 index 00000000000..fcb62242c2d --- /dev/null +++ b/torchtitan_neuron/models/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. diff --git a/torchtitan_neuron/models/deepseek_v3/__init__.py b/torchtitan_neuron/models/deepseek_v3/__init__.py new file mode 100644 index 00000000000..9d17d96baba --- /dev/null +++ b/torchtitan_neuron/models/deepseek_v3/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron-specific DeepSeek V3 model implementations.""" diff --git a/torchtitan_neuron/models/deepseek_v3/config_registry.py b/torchtitan_neuron/models/deepseek_v3/config_registry.py new file mode 100644 index 00000000000..d4a8d73f2ab --- /dev/null +++ b/torchtitan_neuron/models/deepseek_v3/config_registry.py @@ -0,0 +1,71 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron-native DeepSeek-V3 trainer configs.""" + +from torchtitan.components.checkpoint import CheckpointManager +from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss +from torchtitan.components.lr_scheduler import LRSchedulersContainer +from torchtitan.components.metrics import MetricsProcessor +from torchtitan.components.optimizer import default_adamw +from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.distributed.activation_checkpoint import FullAC +from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader +from torchtitan.models.common.config_utils import decoder_vocab_size +from torchtitan.models.deepseek_v3 import model_registry +from torchtitan.trainer import Trainer + +from .nki_moe import to_nki_moe_config + + +def deepseek_v3_16b_neuron_tp4_fsdp16() -> Trainer.Config: + model_spec = model_registry("16B", attn_backend="sdpa") + config = Trainer.Config( + compile=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + loss=ChunkedLossWrapper.Config( + loss_fn=CrossEntropyLoss.Config( + global_vocab_size=decoder_vocab_size(model_spec), + ), + ), + hf_assets_path="./tests/assets/tokenizer", + metrics=MetricsProcessor.Config(log_freq=1), + model_spec=model_spec, + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4_test"), + optimizer=default_adamw(lr=2.2e-4), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=10, + decay_ratio=0.8, + decay_type="cosine", + min_lr_factor=0.1, + ), + training=TrainingConfig( + local_batch_size=4, + global_batch_size=64, + seq_len=4096, + steps=20, + ), + parallelism=ParallelismConfig( + data_parallel_shard_degree=-1, + tensor_parallel_degree=4, + expert_parallel_degree=1, + ), + checkpoint=CheckpointManager.Config( + enable=True, + interval=10, + load_only=True, + ), + activation_checkpoint=FullAC.Config(), + ) + + for layer in config.model_spec.model.layers: + layer.attention.rope.use_prefix_cache_for_positions = True + if layer.moe is not None: + layer.moe = to_nki_moe_config(layer.moe) + + return config diff --git a/torchtitan_neuron/models/deepseek_v3/nki_moe.py b/torchtitan_neuron/models/deepseek_v3/nki_moe.py new file mode 100644 index 00000000000..cf4c634d69d --- /dev/null +++ b/torchtitan_neuron/models/deepseek_v3/nki_moe.py @@ -0,0 +1,658 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""DeepSeek V3 NKI MoE integration for Neuron.""" + +from dataclasses import dataclass +from typing import Literal + +import nki.language as nl + +import spmd_types as spmd +import torch +import torch.nn.functional as F +from nkilib.core.subkernels.find_nonzero_indices import find_nonzero_indices +from nkilib.core.subkernels.indexed_flatten import indexed_flatten +from nkilib.core.utils.common_types import ActFnType, ExpertAffinityScaleMode +from nkilib.experimental.moe.bwd.blockwise_mm_backward import blockwise_mm_bwd +from nkilib.experimental.moe.bwd.moe_bwd_parameters import ( + AffinityOption, + ShardOption, + SkipMode as BwdSkipMode, +) +from nkilib.experimental.moe.forward.bwmm_shard_on_H import ( + blockwise_mm_baseline_shard_hidden, + SkipMode as FwdSkipMode, +) +from torch import nn +from torch.distributed.tensor import DTensor, Partial, Replicate +from torch_neuronx import nki_op +from torch_neuronx.nki_hop import wrap_nki +from torch_neuronx.utils import get_logical_neuron_cores + +from torchtitan.models.common import FeedForward, Linear +from torchtitan.models.common.decoder_sharding import ( + dense_activation_placement, + dense_param_placement, + dense_sequence_parallel_placement, +) +from torchtitan.models.common.moe import GroupedExperts, MoE, TokenChoiceTopKRouter +from torchtitan.protocols.module import Module +from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig + + +BLOCK_SIZE = 512 +LNC = (int(get_logical_neuron_cores()),) + +wrapped_moe_fwd = wrap_nki(blockwise_mm_baseline_shard_hidden) +wrapped_moe_bwd = wrap_nki(blockwise_mm_bwd) +wrapped_find_nonzero = wrap_nki(find_nonzero_indices) +wrapped_indexed_flatten = wrap_nki(indexed_flatten) + + +def _debug_force_load_balance_masked_scores( + scores: torch.Tensor, + top_k: int, + num_experts: int, + route_norm: bool, + route_scale: float, +) -> torch.Tensor: + n_tokens = scores.shape[0] + indices = ( + torch.arange(n_tokens * top_k, device=scores.device, dtype=torch.int64).reshape( + n_tokens, top_k + ) + % num_experts + ) + top_scores = scores.gather(dim=1, index=indices) + if route_norm: + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20) + top_scores = top_scores * route_scale + + masked_scores = torch.zeros_like(scores) + masked_scores.scatter_(1, indices, top_scores) + return masked_scores + + +@nki_op("deepseek_v3_nki_moe::nki_moe_fwd") +def nki_fwd( + hidden_states: torch.Tensor, + masked_scores: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + top_k: int, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + int, + int, + torch.Tensor, +]: + tokens = hidden_states.shape[0] + num_experts = masked_scores.shape[1] + block_size = BLOCK_SIZE + num_blocks = (tokens * top_k + block_size - 1) // block_size + num_experts + + if num_experts == 1: + token_position_to_id = torch.arange( + num_blocks * block_size, + device=hidden_states.device, + dtype=torch.int32, + ) + token_position_to_id = torch.where( + token_position_to_id < tokens, + token_position_to_id, + -1, + ) + block_to_expert = torch.zeros( + (num_blocks, 1), + device=hidden_states.device, + dtype=torch.int32, + ) + else: + indices, counts = wrapped_find_nonzero[LNC](masked_scores.float()) + blocks_per_expert = ((counts + block_size - 1) // block_size).clamp(min=1) + slack = num_blocks - blocks_per_expert.sum() + padded_blocks_per_expert = blocks_per_expert.clone() + padded_blocks_per_expert[-1] += slack + row_offsets = ( + torch.cumsum(padded_blocks_per_expert, dim=0) - padded_blocks_per_expert + ) + expert_ids = torch.arange( + num_experts, device=hidden_states.device, dtype=torch.int32 + ) + block_to_expert = torch.repeat_interleave( + expert_ids, + padded_blocks_per_expert, + output_size=num_blocks, + ).unsqueeze(1) + token_position_to_id = wrapped_indexed_flatten[LNC]( + indices, + block_size, + num_blocks * block_size, + row_offsets, + padding_val=-1, + ) + + affinities = masked_scores.reshape(-1, 1).clone() + gate_up_proj_weight = torch.stack( + (w1.transpose(-2, -1), w3.transpose(-2, -1)), + dim=2, + ) + down_proj_weight = w2.transpose(-2, -1) + hidden_dim_tp = down_proj_weight.shape[1] + gate_up_activations = torch.empty( + (num_blocks, 2, hidden_dim_tp, block_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + down_activations = None + output, gate_up_activations = wrapped_moe_fwd[LNC]( + hidden_states, + affinities, + gate_up_proj_weight, + down_proj_weight, + block_size, + token_position_to_id, + block_to_expert, + gate_up_activations, + down_activations, + skip_dma=FwdSkipMode(skip_token=True), + is_tensor_update_accumulating=True, + expert_affinities_scaling_mode=ExpertAffinityScaleMode.POST_SCALE, + # The kernel sums the top-k expert outputs in bf16 unless told otherwise. + accum_dtype=nl.float32, + ) + + return ( + affinities, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + tokens, + num_experts, + output, + ) + + +@nki_op( + "deepseek_v3_nki_moe::nki_moe_bwd", + mutates_args={ + "hidden_states_grad_out", + "expert_affinities_grad_out", + "gate_up_proj_weight_grad_out", + "down_proj_weight_grad_out", + }, +) +def nki_bwd( + hidden_states: torch.Tensor, + affinities: torch.Tensor, + gate_up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + gate_up_activations: torch.Tensor, + down_activations: torch.Tensor, + token_position_to_id: torch.Tensor, + block_to_expert: torch.Tensor, + grad_output: torch.Tensor, + hidden_states_grad_out: torch.Tensor, + expert_affinities_grad_out: torch.Tensor, + gate_up_proj_weight_grad_out: torch.Tensor, + down_proj_weight_grad_out: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return wrapped_moe_bwd[LNC]( + hidden_states, + affinities, + gate_up_proj_weight, + down_proj_weight, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + grad_output, + BLOCK_SIZE, + skip_dma=BwdSkipMode(skip_token=True), + skip_grad_initialization=True, + activation_type=ActFnType.SiLU, + affinity_option=AffinityOption.AFFINITY_ON_I, + shard_option=ShardOption.SHARD_ON_HIDDEN, + bias=False, + hidden_states_grad_out=hidden_states_grad_out, + expert_affinities_masked_grad_out=expert_affinities_grad_out, + gate_up_proj_weight_grad_out=gate_up_proj_weight_grad_out, + down_proj_weight_grad_out=down_proj_weight_grad_out, + ) + + +def _setup_context(ctx, inputs, output): + hidden_states, masked_scores, w1, w2, w3, _ = inputs + ( + affinities, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + tokens, + num_experts, + _, + ) = output + ctx.save_for_backward( + hidden_states, + affinities, + w1, + w2, + w3, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + ) + ctx.tokens = tokens + ctx.num_experts = num_experts + ctx.scores_dtype = masked_scores.dtype + + +def _backward(ctx, *grad_outputs): + grad_output = grad_outputs[-1] + ( + hidden_states, + affinities, + w1, + w2, + w3, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + ) = ctx.saved_tensors + + gate_up_proj_weight = torch.stack( + (w1.transpose(-2, -1), w3.transpose(-2, -1)), + dim=2, + ) + down_proj_weight = w2.transpose(-2, -1) + + hidden_states_grad_out = torch.zeros_like(hidden_states, dtype=torch.float32) + expert_affinities_grad_out = torch.zeros_like(affinities) + gate_up_proj_weight_grad_out = torch.zeros_like( + gate_up_proj_weight, dtype=torch.float32 + ) + down_proj_weight_grad_out = torch.zeros_like(down_proj_weight, dtype=torch.float32) + + ( + hidden_states_grad, + affinities_grad, + gate_up_proj_weight_grad, + down_proj_weight_grad, + ) = nki_bwd( + hidden_states, + affinities, + gate_up_proj_weight, + down_proj_weight, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + grad_output[: ctx.tokens], + hidden_states_grad_out, + expert_affinities_grad_out, + gate_up_proj_weight_grad_out, + down_proj_weight_grad_out, + ) + + hidden_states_grad = hidden_states_grad.to(hidden_states.dtype) + gate_up_proj_weight_grad = gate_up_proj_weight_grad.to(gate_up_proj_weight.dtype) + down_proj_weight_grad = down_proj_weight_grad.to(down_proj_weight.dtype) + scores_grad = affinities_grad.reshape(-1, ctx.num_experts)[: ctx.tokens].to( + ctx.scores_dtype + ) + w1_grad, w3_grad = gate_up_proj_weight_grad.unbind(dim=2) + + return ( + hidden_states_grad[: ctx.tokens], + scores_grad, + w1_grad.transpose(-2, -1), + down_proj_weight_grad.transpose(-2, -1), + w3_grad.transpose(-2, -1), + None, + ) + + +nki_fwd.register_autograd(_backward, setup_context=_setup_context) + + +def _local(t): + """Plain local view of a (Replicate) DTensor buffer/param for the local MoE region.""" + return t.to_local() if isinstance(t, DTensor) else t + + +def _localize_gate_output(scores: torch.Tensor) -> torch.Tensor: + """Localize the router output with a Replicate local gradient. + + The experts' from_local(Replicate) input wrap already all-reduces the per-rank score + gradients; a Partial label here would sum them again (gate gradient TP x too large).""" + if isinstance(scores, DTensor): + return scores.to_local( + grad_placements=tuple(Replicate() for _ in scores.placements) + ) + return scores + + +class MaskedScoresTopKRouter(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + gate: Linear.Config + top_k: int + score_func: Literal["softmax", "sigmoid"] + route_norm: bool = False + route_scale: float = 1.0 + _debug_force_load_balance: bool = False + + def __init__(self, config: Config): + super().__init__() + self.gate = config.gate.build() + self.top_k = config.top_k + self.score_func = config.score_func + self.route_norm = config.route_norm + self.route_scale = config.route_scale + self._debug_force_load_balance = config._debug_force_load_balance + + def forward( + self, + x: torch.Tensor, + expert_bias: torch.Tensor | None = None, + ) -> torch.Tensor: + with torch.autocast(device_type=x.device.type, dtype=torch.float32): + scores = _localize_gate_output(self.gate(x)) + if self.score_func == "softmax": + scores = F.softmax(scores, dim=-1) + elif self.score_func == "sigmoid": + scores = torch.sigmoid(scores) + else: + raise ValueError(f"Unsupported score function: {self.score_func}") + + scores_for_choice = scores if expert_bias is None else scores + expert_bias + top_indices = torch.topk( + scores_for_choice, + k=self.top_k, + dim=-1, + sorted=False, + ).indices + mask = torch.zeros_like(scores_for_choice, dtype=torch.bool) + mask.scatter_(-1, top_indices, True) + masked_scores = scores * mask + if self.route_norm: + masked_scores = masked_scores / ( + masked_scores.sum(dim=-1, keepdim=True) + 1e-20 + ) + return masked_scores * self.route_scale + + +class NKIGroupedExperts(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + hidden_dim: int + num_experts: int + top_k: int + token_dispatcher: object | None = None + + def __init__(self, config: Config): + super().__init__() + self.num_experts = config.num_experts + self.top_k = config.top_k + self.w1 = nn.Parameter( + torch.empty(config.num_experts, config.hidden_dim, config.dim) + ) + self.w2 = nn.Parameter( + torch.empty(config.num_experts, config.dim, config.hidden_dim) + ) + self.w3 = nn.Parameter( + torch.empty(config.num_experts, config.hidden_dim, config.dim) + ) + + def forward( + self, + hidden_states: torch.Tensor, + masked_scores: torch.Tensor, + ) -> torch.Tensor: + w1 = self.w1 + w2 = self.w2 + w3 = self.w3 + if isinstance(w1, DTensor): + w1 = w1.to_local() + w2 = w2.to_local() + w3 = w3.to_local() + (*_, tokens, _, output) = nki_fwd( + hidden_states, + masked_scores, + w1, + w2, + w3, + self.top_k, + ) + return output[:tokens] + + +class NKIMoE(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + num_experts: int + router: MaskedScoresTopKRouter.Config + experts: NKIGroupedExperts.Config + load_balance_coeff: float | None = 1e-3 + shared_experts: FeedForward.Config | None = None + + def set_sharding_config(self, *, enable_sp: bool, enable_ep: bool) -> None: + if enable_ep: + raise NotImplementedError( + "DeepSeek V3 NKI MoE supports TP-only experts" + ) + + activation = ( + dense_sequence_parallel_placement() + if enable_sp + else dense_activation_placement(tp=spmd.I) + ) + replicated = dense_activation_placement(tp=spmd.R) + partial = dense_activation_placement(tp=spmd.P) + self.sharding_config = ShardingConfig( + state_shardings={ + "expert_bias_E": dense_param_placement(tp=spmd.R), + "tokens_per_expert_E": dense_param_placement(tp=spmd.R), + }, + in_src_shardings={"x": activation}, + in_dst_shardings={"x": replicated}, + out_src_shardings=partial, + out_dst_shardings=activation, + ) + self.router.gate.sharding_config = ShardingConfig( + state_shardings={ + "weight": dense_param_placement(tp=spmd.R), + "bias": dense_param_placement(tp=spmd.R), + }, + in_src_shardings={"input": replicated}, + out_src_shardings=replicated, + ) + self.experts.sharding_config = ShardingConfig( + state_shardings={ + "w1": dense_param_placement(tp=spmd.S(1)), + "w2": dense_param_placement(tp=spmd.S(2)), + "w3": dense_param_placement(tp=spmd.S(1)), + }, + in_src_shardings={ + "hidden_states": replicated, + "masked_scores": replicated, + }, + in_dst_shardings={ + "hidden_states": replicated, + "masked_scores": replicated, + }, + out_src_shardings=partial, + out_dst_shardings=partial, + local_map=LocalMapConfig(in_grad_placements=(partial, partial)), + ) + if self.shared_experts is not None: + shared_hidden = dense_activation_placement(tp=spmd.S(1), cp=spmd.R) + self.shared_experts.sharding_config = ShardingConfig( + in_src_shardings={"x": replicated}, + in_dst_shardings={"x": replicated}, + ) + self.shared_experts.w1.sharding_config = ShardingConfig( + state_shardings={ + "weight": dense_param_placement(tp=spmd.S(0)), + "bias": dense_param_placement(tp=spmd.S(0)), + }, + in_src_shardings={"input": replicated}, + in_dst_shardings={"input": replicated}, + out_src_shardings=shared_hidden, + out_dst_shardings=shared_hidden, + ) + self.shared_experts.w2.sharding_config = ShardingConfig( + state_shardings={ + "weight": dense_param_placement(tp=spmd.S(1)), + "bias": dense_param_placement(tp=spmd.R), + }, + in_src_shardings={"input": shared_hidden}, + out_src_shardings=partial, + out_dst_shardings=partial, + ) + self.shared_experts.w3.sharding_config = ( + self.shared_experts.w1.sharding_config + ) + + def __init__(self, config: Config): + super().__init__() + self.num_experts = config.num_experts + self.router = config.router.build() + self.experts = config.experts.build() + self.shared_experts = ( + config.shared_experts.build() if config.shared_experts is not None else None + ) + self.load_balance_coeff = config.load_balance_coeff + self.moe_enabled = True + if self.load_balance_coeff is not None: + self.register_buffer( + "expert_bias_E", + torch.zeros(self.num_experts, dtype=torch.float32), + persistent=True, + ) + else: + self.expert_bias_E = None + self.register_buffer( + "tokens_per_expert_E", + torch.zeros(self.num_experts, dtype=torch.float32), + persistent=False, + ) + + def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: + if buffer_device is None: + buffer_device = self.tokens_per_expert_E.device + + with torch.device(buffer_device): + self.tokens_per_expert_E = torch.zeros( + self.num_experts, dtype=torch.float32 + ) + if self.load_balance_coeff is not None: + self.expert_bias_E = torch.zeros(self.num_experts, dtype=torch.float32) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Run router + experts on plain local tensors. Each sub-module's + # ShardingConfig re-wraps its plain input (DTensor.from_local on in_src) + # and that wrap's backward all-reduces its Partial input gradient, so the + # gradient arriving here is already fully reduced: localize with a plain + # to_local() (declaring it Partial would sum it a second time). + spec = None + if isinstance(x, DTensor): + spec = (x.device_mesh, x.placements) + x = x.to_local() + batch, sequence, dim = x.shape + flat = x.reshape(-1, dim) + if self.router._debug_force_load_balance: + with torch.autocast(device_type=flat.device.type, dtype=torch.float32): + raw_scores = _localize_gate_output(self.router.gate(flat)) + if self.router.score_func == "sigmoid": + raw_scores = torch.sigmoid(raw_scores) + else: + raw_scores = F.softmax(raw_scores, dim=-1) + masked_scores = _debug_force_load_balance_masked_scores( + raw_scores, + self.router.top_k, + self.num_experts, + self.router.route_norm, + self.router.route_scale, + ) + else: + masked_scores = _localize_gate_output( + self.router(flat, _local(self.expert_bias_E)) + ) + + with torch.no_grad(): + _local(self.tokens_per_expert_E).add_( + (masked_scores != 0).to(torch.float32).sum(dim=0) + ) + + # Both expert paths return Partial DTensors (their ShardingConfig out_src); + # a plain to_local() labels their gradient Replicate, which is the correct + # gradient placement for a Partial output. + routed_output = _local(self.experts(flat, masked_scores)) + if self.shared_experts is not None: + routed_output = routed_output + _local(self.shared_experts(flat)) + routed_output = routed_output.reshape(batch, sequence, dim) + if spec is not None: + routed_output = DTensor.from_local( + routed_output, + spec[0], + tuple(Partial() for _ in spec[1]), + run_check=False, + ) + return routed_output + + +def to_nki_moe_config(moe: MoE.Config) -> NKIMoE.Config: + if not isinstance(moe.router, TokenChoiceTopKRouter.Config): + raise TypeError( + f"Expected TokenChoiceTopKRouter.Config, got {type(moe.router)!r}" + ) + if not isinstance(moe.experts, GroupedExperts.Config): + raise TypeError(f"Expected GroupedExperts.Config, got {type(moe.experts)!r}") + if ( + moe.router.num_expert_groups is not None + or moe.router.num_limited_groups is not None + ): + raise NotImplementedError( + "DeepSeek V3 NKI MoE does not support group-limited routing" + ) + + param_init = moe.experts.param_init or {} + nki_param_init = { + key: param_init[src] + for key, src in (("w1", "w1_EFD"), ("w2", "w2_EDF"), ("w3", "w3_EFD")) + if param_init.get(src) is not None + } + return NKIMoE.Config( + num_experts=moe.num_experts, + load_balance_coeff=moe.load_balance_coeff, + router=MaskedScoresTopKRouter.Config( + gate=moe.router.gate, + top_k=moe.router.top_k, + score_func=moe.router.score_func, + route_norm=moe.router.route_norm, + route_scale=moe.router.route_scale, + _debug_force_load_balance=moe.router._debug_force_load_balance, + ), + experts=NKIGroupedExperts.Config( + dim=moe.experts.dim, + hidden_dim=moe.experts.hidden_dim, + num_experts=moe.experts.num_experts, + top_k=moe.router.top_k, + param_init=nki_param_init, + ), + shared_experts=moe.shared_experts, + ) diff --git a/torchtitan_neuron/models/gpt_oss/__init__.py b/torchtitan_neuron/models/gpt_oss/__init__.py new file mode 100644 index 00000000000..fcb62242c2d --- /dev/null +++ b/torchtitan_neuron/models/gpt_oss/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. diff --git a/torchtitan_neuron/models/gpt_oss/config_registry.py b/torchtitan_neuron/models/gpt_oss/config_registry.py new file mode 100644 index 00000000000..1e425ae4c96 --- /dev/null +++ b/torchtitan_neuron/models/gpt_oss/config_registry.py @@ -0,0 +1,63 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron-native GPT-OSS trainer configs.""" + +from torchtitan.components.checkpoint import CheckpointManager +from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss +from torchtitan.components.lr_scheduler import LRSchedulersContainer +from torchtitan.components.metrics import MetricsProcessor +from torchtitan.components.optimizer import default_adamw +from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.distributed.activation_checkpoint import FullAC +from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader +from torchtitan.models.common.config_utils import decoder_vocab_size +from torchtitan.models.gpt_oss import model_registry +from torchtitan.trainer import Trainer + +from .nki_moe import to_nki_moe_config + + +def gpt_oss_20b_nki_tp4_fsdp16() -> Trainer.Config: + model_spec = model_registry("20b", attn_backend="nki") + for layer in model_spec.model.layers: + layer.moe = to_nki_moe_config(layer.moe) + # Neuron does not support the RoPE gather; use the prefix-cache path. + layer.attention.rope.use_prefix_cache_for_positions = True + return Trainer.Config( + compile=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + loss=ChunkedLossWrapper.Config( + loss_fn=CrossEntropyLoss.Config( + global_vocab_size=decoder_vocab_size(model_spec), + ), + ), + hf_assets_path="./tests/assets/tokenizer", + metrics=MetricsProcessor.Config(log_freq=1), + model_spec=model_spec, + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4_test"), + optimizer=default_adamw(lr=8e-4), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=10, + decay_ratio=0.8, + decay_type="linear", + min_lr_factor=0.0, + ), + training=TrainingConfig( + local_batch_size=1, + seq_len=4096, + steps=10, + ), + parallelism=ParallelismConfig( + data_parallel_shard_degree=-1, + tensor_parallel_degree=4, + expert_parallel_degree=1, + ), + checkpoint=CheckpointManager.Config(interval=10), + activation_checkpoint=FullAC.Config(), + ) diff --git a/torchtitan_neuron/models/gpt_oss/nki_attention.py b/torchtitan_neuron/models/gpt_oss/nki_attention.py new file mode 100644 index 00000000000..f5761ea0fce --- /dev/null +++ b/torchtitan_neuron/models/gpt_oss/nki_attention.py @@ -0,0 +1,339 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""NKI Flash Attention wrapper for GPT-OSS on Neuron. + +Calls nkilib attention_cte (forward) and attention_bwd (backward) directly, +with native support for attention sinks, sliding window, causal masking, and GQA. + +Uses nki_op + register_autograd for torch.compile compatibility. +""" + +from dataclasses import dataclass + +import spmd_types as spmd +import torch +from nkilib.core.attention.attention_bwd import attention_bwd +from nkilib.core.attention.attention_cte import attention_cte +from torch.distributed.tensor import DTensor +from torch_neuronx import nki_op + +from torch_neuronx.nki_hop import wrap_nki +from torch_neuronx.utils import get_logical_neuron_cores +from torchtitan.models.common.decoder_sharding import ( + dense_activation_placement, + dense_param_placement, +) +from torchtitan.protocols.module import Module +from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig + +wrapped_flash_fwd = wrap_nki(attention_cte) +wrapped_flash_bwd = wrap_nki(attention_bwd) + +LNC = (int(get_logical_neuron_cores()),) +NAMESPACE = "gptoss_nki_attn" + + +def _nki_attn_fwd_impl( + q: torch.Tensor, # (B, S, H, D) + k: torch.Tensor, # (B, S, H_kv, D) + v: torch.Tensor, # (B, S, H_kv, D) + sinks: torch.Tensor, # (H,) + scale: float, + sw: int | None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, S, H, D = q.shape + H_kv = k.shape[2] + + q_nki = q.transpose(1, 2).contiguous().reshape(B * H, S, D) + k_nki = k.transpose(1, 2).contiguous().reshape(B * H_kv, S, D) + v_nki = v.transpose(1, 2).contiguous().reshape(B * H_kv, S, D) + + # When using sliding window, NKI requires scale=1.0; pre-scale Q instead. + if sw is not None: + q_nki = q_nki * scale + kernel_scale = 1.0 + else: + kernel_scale = scale + + sink_nki = sinks.unsqueeze(0).expand(B, -1).reshape(B * H, 1) + + output, neg_max, sum_recip = wrapped_flash_fwd[LNC]( + q_nki, + k_nki, + v_nki, + scale=kernel_scale, + causal_mask=True, + sink=sink_nki, + sliding_window=sw, + tp_q=True, + tp_k=True, + cache_softmax=True, + ) + + lse = -(neg_max + torch.log(sum_recip)) + lse = lse.to(torch.float32).reshape(B, H, *lse.shape[-2:]) + # Output back to (B, S, H, D) + output = output.reshape(B, H, S, D).transpose(1, 2).contiguous() + + return output, lse + + +# --------------------------------------------------------------------------- +# nki_op-based forward +# --------------------------------------------------------------------------- +@nki_op(f"{NAMESPACE}::nki_attn_fwd_full") +def nki_attn_fwd_full( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sinks: torch.Tensor, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + return _nki_attn_fwd_impl(q, k, v, sinks, scale, None) + + +@nki_op(f"{NAMESPACE}::nki_attn_fwd_window_128") +def nki_attn_fwd_window_128( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sinks: torch.Tensor, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + return _nki_attn_fwd_impl(q, k, v, sinks, scale, 128) + + +def _nki_attn_bwd_impl( + q: torch.Tensor, # (B, S, H, D) + k: torch.Tensor, # (B, S, H_kv, D) + v: torch.Tensor, # (B, S, H_kv, D) + out: torch.Tensor, # (B, S, H, D) + grad_out: torch.Tensor, # (B, S, H, D) + lse: torch.Tensor, + sinks: torch.Tensor, + scale: float, + sw: int | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, S, H, D = q.shape + + # (B, S, H, D) → (B, H, D, S) for the kernel's expected bwd layout. + q_t = q.permute(0, 2, 3, 1).contiguous() + k_t = k.permute(0, 2, 3, 1).contiguous() + v_t = v.permute(0, 2, 3, 1).contiguous() + out_t = out.permute(0, 2, 3, 1).contiguous() + grad_out_t = grad_out.permute(0, 2, 3, 1).contiguous() + sink_bwd = sinks.unsqueeze(0).expand(B, -1) + + grad_q_t, grad_k_t, grad_v_t, grad_sinks = wrapped_flash_bwd[LNC]( + q_t, + k_t, + v_t, + out_t, + grad_out_t, + lse, + sinks_ref=sink_bwd, + use_causal_mask=True, + mixed_precision=True, + softmax_scale=scale, + sliding_window=sw, + ) + + # Permute grads back to (B, S, H, D) so they match the forward's input + return ( + grad_q_t.permute(0, 3, 1, 2).contiguous(), + grad_k_t.permute(0, 3, 1, 2).contiguous(), + grad_v_t.permute(0, 3, 1, 2).contiguous(), + grad_sinks.sum(dim=0), + ) + + +# --------------------------------------------------------------------------- +# nki_op-based backward +# --------------------------------------------------------------------------- +@nki_op(f"{NAMESPACE}::nki_attn_bwd") +def nki_attn_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + grad_out: torch.Tensor, + lse: torch.Tensor, + sinks: torch.Tensor, + scale: float, + sliding_window: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + sw = sliding_window if sliding_window != 0 else None + return _nki_attn_bwd_impl(q, k, v, out, grad_out, lse, sinks, scale, sw) + + +# --------------------------------------------------------------------------- +# Autograd wiring +# --------------------------------------------------------------------------- +def _nki_attn_fwd_setup_context(ctx, inputs, output): + q, k, v, sinks, scale = inputs + out, lse = output + ctx.save_for_backward(q, k, v, sinks, out, lse) + ctx.scale = scale + + +def _nki_attn_fwd_full_backward(ctx, *grad_outputs): + # grad_outputs: one per forward output (2); output is at index 0 + grad_out = grad_outputs[0] + q, k, v, sinks, out, lse = ctx.saved_tensors + + grad_q, grad_k, grad_v, grad_sinks = nki_attn_bwd( + q, + k, + v, + out, + grad_out, + lse, + sinks, + ctx.scale, + 0, + ) + + return grad_q, grad_k, grad_v, grad_sinks, None + + +def _nki_attn_fwd_window_128_backward(ctx, *grad_outputs): + # grad_outputs: one per forward output (2); output is at index 0 + grad_out = grad_outputs[0] + q, k, v, sinks, out, lse = ctx.saved_tensors + + grad_q, grad_k, grad_v, grad_sinks = nki_attn_bwd( + q, + k, + v, + out, + grad_out, + lse, + sinks, + ctx.scale, + 128, + ) + + return grad_q, grad_k, grad_v, grad_sinks, None + + +nki_attn_fwd_full.register_autograd( + _nki_attn_fwd_full_backward, + setup_context=_nki_attn_fwd_setup_context, +) +nki_attn_fwd_window_128.register_autograd( + _nki_attn_fwd_window_128_backward, + setup_context=_nki_attn_fwd_setup_context, +) + + +# --------------------------------------------------------------------------- +# Module wrapper (unchanged interface) +# --------------------------------------------------------------------------- +class NKIFlashAttentionWrapper(Module): + """Drop-in inner_attention replacement for GPT-OSS Attention. + + Accepts q/k/v in (B, S, H, D) layout (the natural output of view/cat from + the linear projections). The transpose to head-major (B*H, S, D) layout + expected by the NKI kernel happens inside the @nki_op so dynamo never + sees non-contiguous strided tensors in the compiled HLO graph. + + Handles DTensor unwrap/wrap for TP compatibility. + Sinks, sliding window, and causal masking are handled natively by the NKI kernel. + """ + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + softmax_scale: float + sliding_window: int | None = None + + def set_sharding_config(self) -> None: + """Set the TP-sharded q/k/v/sinks local region for the NKI kernel. + + Heads are sharded on TP (``Shard(2)`` of ``(B, S, H, D)``). Under CP, + k/v are gathered to Replicate and their gradients return Partial. + """ + q_placements = dense_activation_placement(tp=spmd.S(2)) + kv_src_placements = dense_activation_placement(tp=spmd.S(2)) + kv_dst_placements = dense_activation_placement(tp=spmd.S(2), cp=spmd.R) + kv_grad_placements = dense_activation_placement(tp=spmd.S(2), cp=spmd.P) + sinks_placements = dense_param_placement(tp=spmd.S(0)) + self.sharding_config = ShardingConfig( + in_src_shardings={ + "q": q_placements, + "k": kv_src_placements, + "v": kv_src_placements, + "sinks": sinks_placements, + }, + in_dst_shardings={ + "q": q_placements, + "k": kv_dst_placements, + "v": kv_dst_placements, + "sinks": sinks_placements, + }, + out_src_shardings=q_placements, + local_map=LocalMapConfig( + in_grad_placements=( + q_placements, + kv_grad_placements, + kv_grad_placements, + sinks_placements, + ), + ), + ) + + uses_attention_sinks_natively = True + + def __init__(self, config: Config): + super().__init__() + if config.sliding_window not in (None, 128): + raise ValueError( + "GPT-OSS NKI attention only supports sliding_window=None or 128, " + f"got {config.sliding_window!r}" + ) + self.softmax_scale = config.softmax_scale + # Use 0 to represent None (nki_op needs concrete int, not Optional) + self.sliding_window_int = ( + config.sliding_window if config.sliding_window is not None else 0 + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sinks: torch.Tensor, + ) -> torch.Tensor: + mesh = None + placements = None + if isinstance(q, DTensor): + mesh = q.device_mesh + placements = q.placements + q = q.to_local() + k = k.to_local() + v = v.to_local() + if isinstance(sinks, DTensor): + sinks = sinks.to_local() + + if self.sliding_window_int == 128: + (*_saved, out, _lse) = nki_attn_fwd_window_128( + q, + k, + v, + sinks, + self.softmax_scale, + ) + else: + (*_saved, out, _lse) = nki_attn_fwd_full( + q, + k, + v, + sinks, + self.softmax_scale, + ) + + if mesh is not None: + out = DTensor.from_local(out, mesh, placements) + return out diff --git a/torchtitan_neuron/models/gpt_oss/nki_moe.py b/torchtitan_neuron/models/gpt_oss/nki_moe.py new file mode 100644 index 00000000000..61ca9d68a7d --- /dev/null +++ b/torchtitan_neuron/models/gpt_oss/nki_moe.py @@ -0,0 +1,760 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""NKI MoE integration for GPT-OSS on Neuron. + +GPT-OSS-specific pieces: +- SwiGLU with clamp (gate/up clamped at ±swiglu_limit; default 7.0) +- Fused (gate+up) weight layout: mlp1_weight (E, H, 2, I), mlp2_weight (E, I, H) +- +1 baked into up_bias (GPT-OSS init convention) +- ActFnType.Swish + +This file intentionally keeps GPT-OSS NKI helpers local instead of refactoring +Qwen3 NKI or deprecated override paths. +""" + +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Literal + +import nki.language as nl + +import spmd_types as spmd +import torch +import torch.distributed._functional_collectives +import torch.nn.functional as F +from nkilib.core.moe.moe_cte.moe_cte import ( + moe_cte, + MoECTEImplementation, + MoECTESpec, + ShardOnIConfig, +) + +from nkilib.core.subkernels.find_nonzero_indices import find_nonzero_indices +from nkilib.core.subkernels.indexed_flatten import indexed_flatten +from nkilib.core.utils.common_types import ActFnType, ExpertAffinityScaleMode +from nkilib.experimental.moe.bwd.blockwise_mm_backward import blockwise_mm_bwd +from nkilib.experimental.moe.bwd.moe_bwd_parameters import ( + AffinityOption, + ClampLimits, + ShardOption, + SkipMode, +) +from torch import nn +from torch.distributed.tensor import DTensor, Partial +from torch_neuronx import nki_op + +from torch_neuronx.nki_hop import wrap_nki +from torch_neuronx.utils import get_logical_neuron_cores + +from torchtitan.models.common import Linear +from torchtitan.models.common.decoder_sharding import ( + dense_activation_placement, + dense_param_placement, + dense_sequence_parallel_placement, +) +from torchtitan.models.common.moe import TokenChoiceTopKRouter +from torchtitan.models.gpt_oss.moe import GptOssGroupedExperts, GptOssMoE +from torchtitan.protocols.module import Module +from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig + + +# --------------------------------------------------------------------------- +# GPT-OSS-specific kernel handles & constants +# --------------------------------------------------------------------------- +LNC = (int(get_logical_neuron_cores()),) + +wrapped_moe_fwd = wrap_nki(moe_cte) +wrapped_moe_bwd = wrap_nki(blockwise_mm_bwd) +wrapped_find_nonzero = wrap_nki(find_nonzero_indices) +wrapped_indexed_flatten = wrap_nki(indexed_flatten) + +NAMESPACE = "gptoss_nki_moe" + +# The installed nkilib MoE kernel requires the hidden (H) and intermediate (I) +# dims to be multiples of this. GPT-OSS dims (2880) are not, so the local +# tensors are zero-padded up to the nearest multiple just for the kernel call. +KERNEL_DIM_MULTIPLE = 512 + + +def _round_up(x: int, multiple: int) -> int: + return -(-x // multiple) * multiple + + +def _compute_n_max(tokens: int, top_k: int, block_size: int, num_experts: int) -> int: + return (tokens * top_k + block_size - 1) // block_size + (num_experts - 1) + + +def _debug_force_load_balance_masked_scores( + scores: torch.Tensor, + top_k: int, + num_experts: int, + route_norm: bool, + route_scale: float, +) -> torch.Tensor: + n_tokens = scores.shape[0] + indices = ( + torch.arange(n_tokens * top_k, device=scores.device, dtype=torch.int64).reshape( + n_tokens, top_k + ) + % num_experts + ) + top_scores = scores.gather(dim=1, index=indices) + if route_norm: + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20) + top_scores = top_scores * route_scale + + masked_scores = torch.zeros_like(scores) + masked_scores.scatter_(1, indices, top_scores) + return masked_scores + + +def _local(t): + """Plain local view of a (Replicate) DTensor buffer/param for the local MoE region.""" + return t.to_local() if isinstance(t, DTensor) else t + + +class _SumGradOverGroup(torch.autograd.Function): + """Identity forward; backward all-reduces (sum) the gradient over ``group``.""" + + @staticmethod + def forward(ctx, x, group): + ctx.group = group + return x.view_as(x) + + @staticmethod + def backward(ctx, grad): + if ctx.group is not None: + grad = torch.distributed._functional_collectives.all_reduce( + grad, "sum", ctx.group + ) + return grad, None + + +def _local_gate(gate: nn.Module, x: torch.Tensor, tp_group) -> torch.Tensor: + """Router gate on local tensors (weight/bias are Replicate on TP). + + Runs outside DTensor so the weight gradient and the input gradient see one + explicit TP reduction of the incoming scores gradient (each TP rank's experts + shard returns a per-rank partial). ``tp_group=None`` skips the reduction. + """ + w = gate.weight.to_local() if isinstance(gate.weight, DTensor) else gate.weight + b = gate.bias + if b is not None and isinstance(b, DTensor): + b = b.to_local() + # Reduce on the INPUT edge only (what mainline's NoParallel(Partial) on a + # DTensor plan effectively does): the router-path gradient into x is summed + # over TP, while the gate weight/bias gradient is taken from the local + # scores gradient. + return F.linear(_SumGradOverGroup.apply(x, tp_group), w, b) + + +class MaskedScoresTopKRouter(Module): + """GPT-OSS NKI router that emits a dense ``(tokens, experts)`` score mask.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + gate: Linear.Config + top_k: int + score_func: Literal["softmax", "sigmoid"] = "softmax" + route_norm: bool = True + route_scale: float = 1.0 + _debug_force_load_balance: bool = False + + def __init__(self, config: Config): + super().__init__() + self.gate = config.gate.build() + self.top_k = config.top_k + self.num_experts = config.gate.out_features + self.score_func = config.score_func + self.route_norm = config.route_norm + self.route_scale = config.route_scale + self._debug_force_load_balance = config._debug_force_load_balance + + def forward( + self, + x: torch.Tensor, + expert_bias: torch.Tensor | None = None, + tp_group=None, + ) -> torch.Tensor: + with torch.autocast(device_type=x.device.type, dtype=torch.float32): + scores = _local_gate(self.gate, x, tp_group) + if self.score_func == "softmax": + scores = F.softmax(scores, dim=-1) + elif self.score_func == "sigmoid": + scores = torch.sigmoid(scores) + else: + raise ValueError(f"Unsupported score function: {self.score_func}") + + if self._debug_force_load_balance: + return _debug_force_load_balance_masked_scores( + scores, + self.top_k, + self.num_experts, + self.route_norm, + self.route_scale, + ) + + scores_for_choice = scores if expert_bias is None else scores + expert_bias + top_indices = torch.topk( + scores_for_choice, k=self.top_k, dim=-1, sorted=False + ).indices + mask = torch.zeros_like(scores_for_choice, dtype=torch.bool) + mask.scatter_(-1, top_indices, True) + masked_scores = scores * mask + if self.route_norm: + masked_scores = masked_scores / ( + masked_scores.sum(dim=-1, keepdim=True) + 1e-20 + ) + return masked_scores * self.route_scale + + +# Backward-kernel blocking hyperparameters. +# --------------------------------------------------------------------------- +# nki_op-based forward — compile-compatible +# --------------------------------------------------------------------------- +@nki_op(f"{NAMESPACE}::nki_moe_fwd") +def nki_moe_fwd( + hidden_states: torch.Tensor, # (T, H) + masked_scores: torch.Tensor, # (T, E) float32, top_k nonzero per row + mlp1_weight: torch.Tensor, # (E, H, 2, I) + mlp1_bias: torch.Tensor, # (E, 2, I) + mlp2_weight: torch.Tensor, # (E, I, H) + mlp2_bias: torch.Tensor, # (E, H) + top_k: int, + tp_degree: int = 1, + block_size: int = 256, + swiglu_limit: float = 7.0, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int, torch.Tensor +]: + T, H = hidden_states.shape + E = masked_scores.shape[1] + B = block_size + N_max = _compute_n_max(T, top_k, B, E) + + # --- Index calculation (NKI kernels) --- + indices, counts = wrapped_find_nonzero[LNC](masked_scores.float()) + blocks_per_expert = (counts + B - 1) // B + + # Tail slack: pad last expert so sum == N_max + slack = N_max - blocks_per_expert.sum() + bpe_padded = blocks_per_expert.clone() + # Use explicit positive index (E-1) instead of [-1] to avoid + # torch-mlir slice_scatter lowering bug with negative indices + last_idx = E - 1 + bpe_padded[last_idx] = bpe_padded[last_idx] + slack + + row_offsets = torch.cumsum(bpe_padded, 0) - bpe_padded + + expert_ids = torch.arange(E, device=hidden_states.device, dtype=torch.int32) + block_to_expert = torch.repeat_interleave( + expert_ids, bpe_padded, output_size=N_max + ).unsqueeze(1) + + # padding_val=-1: indexed_flatten's LNC2 merge uses signed nl.maximum; sentinel + # must be < every real value (real token ids are in [0, T)). + token_position_to_id = wrapped_indexed_flatten[LNC]( + indices, + B, + N_max * B, + row_offsets, + padding_val=-1, + ) + + # With skip_dma=SkipMode(skip_token=True), the MoE kernel uses oob.skip + # on token-indirect DMAs — sentinel index -1 is OOB and the DMA is + # suppressed. No padding needed. + affinities_fwd = masked_scores.to(hidden_states.dtype).reshape(-1, 1) + # clone: a custom-op return must not alias an input (masked_scores); reshape + # returns a view, which the eager torch._library aliasing check rejects. + affinities = masked_scores.reshape(-1, 1).clone() + + if tp_degree > 1: + mlp2_bias = mlp2_bias / tp_degree + + # --- Forward kernel --- + spec = MoECTESpec( + implementation=MoECTEImplementation.shard_on_i, + shard_on_I=ShardOnIConfig( + checkpoint_activation=True, + expert_affinity_multiply_on_I=True, + ), + ) + + # With expert_affinity_multiply_on_I=True, the kernel returns + # (output, gate_up_activations_T) — NO down_activations. + output, gup_act_T = wrapped_moe_fwd[LNC]( + hidden_states, + affinities_fwd, # bf16 to fwd kernel; fp32 is saved for bwd via the autograd output tuple + mlp1_weight, + mlp2_weight, + token_position_to_id, + block_to_expert, + B, + spec, + gate_and_up_proj_bias=mlp1_bias, + down_proj_bias=mlp2_bias, + activation_function=ActFnType.Swish, + expert_affinities_scaling_mode=ExpertAffinityScaleMode.POST_SCALE, + skip_dma=SkipMode(skip_token=True), + is_tensor_update_accumulating=True, + gate_clamp_upper_limit=swiglu_limit, + # +1 is baked into the up-bias, which the kernel clamps BEFORE the + # swiglu multiply. GPT-OSS clamps the raw up then adds +1, so shift + # the up-clamp window by +1: clamp(up+1, -L+1, L+1) == clamp(up, -L, L) + 1. + up_clamp_lower_limit=-swiglu_limit + 1.0, + up_clamp_upper_limit=swiglu_limit + 1.0, + accumulation_dtype=nl.float32, + ) + + return (affinities, gup_act_T, token_position_to_id, block_to_expert, T, E, output) + + +# --------------------------------------------------------------------------- +# nki_op-based backward — compile-compatible +# --------------------------------------------------------------------------- +@nki_op(f"{NAMESPACE}::nki_moe_bwd") +def nki_moe_bwd( + hidden_states: torch.Tensor, + affinities: torch.Tensor, + mlp1_weight: torch.Tensor, + mlp2_weight: torch.Tensor, + gup_act_T: torch.Tensor, + token_position_to_id: torch.Tensor, + block_to_expert: torch.Tensor, + grad_out: torch.Tensor, + B: int, + swiglu_limit: float = 7.0, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + clamp_limits = ClampLimits( + non_linear_clamp_upper_limit=swiglu_limit, + linear_clamp_lower_limit=-swiglu_limit + 1.0, + linear_clamp_upper_limit=swiglu_limit + 1.0, + ) + + return wrapped_moe_bwd[LNC]( + hidden_states, + affinities, + mlp1_weight, + mlp2_weight, + gup_act_T, + None, # down_proj_act_checkpoint — not used with AFFINITY_ON_I + token_position_to_id, + block_to_expert, + grad_out, + B, + activation_type=ActFnType.Swish, + affinity_option=AffinityOption.AFFINITY_ON_I, + shard_option=ShardOption.SHARD_ON_HIDDEN, + clamp_limits=clamp_limits, + bias=True, + skip_dma=SkipMode(skip_token=True), + accumulation_dtype=nl.float32, + ) + + +# --------------------------------------------------------------------------- +# Autograd wiring via register_autograd +# --------------------------------------------------------------------------- +def _nki_moe_fwd_setup_context(ctx, inputs, output): + ( + hidden_states, + masked_scores, + mlp1_weight, + mlp1_bias, + mlp2_weight, + mlp2_bias, + top_k, + tp_degree, + block_size, + swiglu_limit, + ) = inputs + ( + affinities, + gup_act_T, + token_position_to_id, + block_to_expert, + T, + E, + _output, + ) = output + ctx.save_for_backward( + hidden_states, + affinities, + mlp1_weight, + mlp2_weight, + gup_act_T, + token_position_to_id, + block_to_expert, + mlp2_bias, + ) + ctx.T = T + ctx.E = E + ctx.block_size = block_size + ctx.swiglu_limit = swiglu_limit + ctx.scores_dtype = masked_scores.dtype + ctx.tp_degree = tp_degree + + +def _nki_moe_fwd_backward(ctx, *grad_outputs): + # grad_outputs has one grad per forward output (10 total); only the last (output) matters + grad_out_full = grad_outputs[-1] + ( + hidden_states, + affinities, + mlp1_weight, + mlp2_weight, + gup_act_T, + token_position_to_id, + block_to_expert, + mlp2_bias, + ) = ctx.saved_tensors + T, E = ctx.T, ctx.E + B = ctx.block_size + + grad_out = grad_out_full[:T] + + ( + hidden_grad, + affinities_grad, + mlp1_weight_grad, + mlp2_weight_grad, + mlp1_bias_grad, + mlp2_bias_grad, + ) = nki_moe_bwd( + hidden_states, + affinities, + mlp1_weight, + mlp2_weight, + gup_act_T, + token_position_to_id, + block_to_expert, + grad_out, + B, + ctx.swiglu_limit, + ) + + scores_grad = affinities_grad.reshape(-1, E)[:T] + + # Kernel's AFFINITY_ON_I affinity grad omits the down-bias term, but the + # forward scales the full down output (incl. bias) by affinity. Add it back. + eff_bias = mlp2_bias.to(torch.float32) + if ctx.tp_degree > 1: + eff_bias = eff_bias / ctx.tp_degree + bias_corr = grad_out.to(torch.float32) @ eff_bias.t() # (T, E) + routed_mask = (affinities.reshape(-1, E)[:T] != 0).to(bias_corr.dtype) + scores_grad = (scores_grad.to(torch.float32) + bias_corr * routed_mask).to( + ctx.scores_dtype + ) + + return ( + hidden_grad[:T], # hidden_states + scores_grad, # masked_scores + mlp1_weight_grad, # mlp1_weight + mlp1_bias_grad, # mlp1_bias + mlp2_weight_grad, # mlp2_weight + mlp2_bias_grad, # mlp2_bias + None, # top_k + None, # tp_degree + None, # block_size + None, # swiglu_limit + ) + + +nki_moe_fwd.register_autograd( + _nki_moe_fwd_backward, setup_context=_nki_moe_fwd_setup_context +) + + +# --------------------------------------------------------------------------- +# Module classes — GPT-OSS subclasses of common base classes +# --------------------------------------------------------------------------- +class NKIGroupedExperts(Module): + """GPT-OSS specific blockwise experts. + + Weight layout (fused gate+up): + mlp1_weight: (E, H, 2, I) + mlp1_bias: (E, 2, I) + mlp2_weight: (E, I, H) + mlp2_bias: (E, H) + + Quirks: + - SwiGLU clamp ±7.0 (handled by the kernel via ``gate_clamp_*`` kwargs) + - ``+1`` baked into the up-bias at init + - 2880→3072 hidden-dim padding so the installed nkilib kernel accepts H + - Calls the GPT-OSS-specific ``nki_moe_fwd`` op + """ + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + hidden_dim: int + num_experts: int + top_k: int + block_size: int = 256 + swiglu_limit: float = 7.0 + tp_degree: int = 1 + + def __init__(self, config: Config): + super().__init__() + self.dim = config.dim + self.hidden_dim = config.hidden_dim + self.num_experts = config.num_experts + self.top_k = config.top_k + self.block_size = config.block_size + self.swiglu_limit = config.swiglu_limit + self._tp_degree = config.tp_degree + + E = self.num_experts + H = self.dim + I = self.hidden_dim + self.mlp1_weight = nn.Parameter(torch.empty(E, H, 2, I)) + self.mlp1_bias = nn.Parameter(torch.empty(E, 2, I)) + self.mlp2_weight = nn.Parameter(torch.empty(E, I, H)) + self.mlp2_bias = nn.Parameter(torch.empty(E, H)) + + def forward(self, hidden_states, masked_scores): + """Triggers FSDP all-gather. Runs the GPT-OSS NKI MoE forward op.""" + if isinstance(self.mlp1_weight, DTensor): + mlp1_w = self.mlp1_weight.to_local() + mlp1_b = self.mlp1_bias.to_local() + mlp2_w = self.mlp2_weight.to_local() + mlp2_b = self.mlp2_bias.to_local() + else: + mlp1_w = self.mlp1_weight + mlp1_b = self.mlp1_bias + mlp2_w = self.mlp2_weight + mlp2_b = self.mlp2_bias + + tp_degree = self._tp_degree + if isinstance(self.mlp1_weight, DTensor): + mesh_dim_names = self.mlp1_weight.device_mesh.mesh_dim_names + if "tp" in mesh_dim_names: + tp_dim_idx = mesh_dim_names.index("tp") + tp_degree = self.mlp1_weight.device_mesh.size(tp_dim_idx) + + # Zero-pad H and I up to a kernel-legal multiple of KERNEL_DIM_MULTIPLE. + H_orig = hidden_states.shape[-1] + I_orig = mlp1_w.shape[-1] # I per TP rank + tp = tp_degree if tp_degree > 0 else 1 + H_target = _round_up(H_orig, KERNEL_DIM_MULTIPLE) + I_target = _round_up(I_orig * tp, KERNEL_DIM_MULTIPLE) // tp + + needs_pad = (H_orig != H_target) or (I_orig != I_target) + if needs_pad: + H_pad = H_target - H_orig + I_pad = I_target - I_orig + hidden_states = torch.nn.functional.pad(hidden_states, (0, H_pad)) + mlp1_w = torch.nn.functional.pad(mlp1_w, (0, I_pad, 0, 0, 0, H_pad)) + mlp1_b = torch.nn.functional.pad(mlp1_b, (0, I_pad)) + mlp2_w = torch.nn.functional.pad(mlp2_w, (0, H_pad, 0, I_pad)) + mlp2_b = torch.nn.functional.pad(mlp2_b, (0, H_pad)) + + (*_saved, T, _E, output) = nki_moe_fwd( + hidden_states, + masked_scores, + mlp1_w, + mlp1_b, + mlp2_w, + mlp2_b, + self.top_k, + tp_degree, + self.block_size, + self.swiglu_limit, + ) + return output[:T, :H_orig] + + +class NKIMoE(Module): + """GPT-OSS NKI MoE using shared masked-score routing helpers.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + num_experts: int + router: MaskedScoresTopKRouter.Config + experts: NKIGroupedExperts.Config + load_balance_coeff: float | None = 1e-3 + shared_experts: None = None + + def set_sharding_config(self, *, enable_sp: bool, enable_ep: bool) -> None: + if enable_ep: + raise NotImplementedError("GPT-OSS NKI MoE supports TP-only experts") + + activation = ( + dense_sequence_parallel_placement() + if enable_sp + else dense_activation_placement(tp=spmd.I) + ) + replicated = dense_activation_placement(tp=spmd.R) + partial = dense_activation_placement(tp=spmd.P) + + self.sharding_config = ShardingConfig( + state_shardings={ + "expert_bias_E": dense_param_placement(tp=spmd.R), + "tokens_per_expert_E": dense_param_placement(tp=spmd.R), + }, + in_src_shardings={"x": activation}, + in_dst_shardings={"x": replicated}, + out_src_shardings=partial, + out_dst_shardings=activation, + ) + self.router.gate.sharding_config = ShardingConfig( + state_shardings={ + "weight": dense_param_placement(tp=spmd.R), + "bias": dense_param_placement(tp=spmd.R), + }, + in_src_shardings={"input": replicated}, + out_src_shardings=replicated, + ) + self.experts.sharding_config = ShardingConfig( + state_shardings={ + "mlp1_weight": dense_param_placement(tp=spmd.S(3)), + "mlp1_bias": dense_param_placement(tp=spmd.S(2)), + "mlp2_weight": dense_param_placement(tp=spmd.S(1)), + "mlp2_bias": dense_param_placement(tp=spmd.R), + }, + in_src_shardings={ + "hidden_states": replicated, + "masked_scores": replicated, + }, + in_dst_shardings={ + "hidden_states": replicated, + "masked_scores": replicated, + }, + out_src_shardings=partial, + out_dst_shardings=partial, + local_map=LocalMapConfig( + in_grad_placements=(partial, partial), + ), + ) + + def __init__(self, config: Config): + super().__init__() + self.num_experts = config.num_experts + self.router = config.router.build() + self.experts = config.experts.build() + self.top_k = config.router.top_k + self.load_balance_coeff = config.load_balance_coeff + self.shared_experts = None + self.moe_enabled = True + + if self.load_balance_coeff is not None: + self.register_buffer( + "expert_bias_E", + torch.zeros(self.num_experts, dtype=torch.float32), + persistent=True, + ) + else: + self.expert_bias_E = None + self.register_buffer( + "tokens_per_expert_E", + torch.zeros(self.num_experts, dtype=torch.float32), + persistent=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Run router + experts on plain local tensors. Each sub-module's + # ShardingConfig re-wraps its plain input (DTensor.from_local on in_src) + # and that wrap's backward all-reduces its Partial input gradient, so the + # gradient arriving here is already fully reduced: localize with a plain + # to_local(). Declaring it Partial sums it a second time (4x ffn_norm grad). + spec = None + tp_group = None + if isinstance(x, DTensor): + assert x.device_mesh.ndim == 1 + assert x.device_mesh.mesh_dim_names == ("tp",) + spec = (x.device_mesh, x.placements) + tp_group = x.device_mesh.get_group("tp") + x = x.to_local() + + batch, sequence, dim = x.shape + flat = x.reshape(-1, dim) + # router.forward already localizes the gate output (Partial grad label); + # its return value is a plain local tensor here. + masked_scores = self.router(flat, _local(self.expert_bias_E), tp_group) + + with torch.no_grad(): + _local(self.tokens_per_expert_E).add_( + (masked_scores != 0).to(torch.float32).sum(dim=0) + ) + + output = self.experts(flat, masked_scores) + if isinstance(output, DTensor): + output = output.to_local() + output = output.reshape(batch, sequence, dim) + if spec is not None: + output = DTensor.from_local( + output, spec[0], tuple(Partial() for _ in spec[1]), run_check=False + ) + return output + + def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: + if buffer_device is None: + buffer_device = self.tokens_per_expert_E.device + with torch.device(buffer_device): + self.tokens_per_expert_E = torch.zeros( + self.num_experts, dtype=torch.float32 + ) + if self.load_balance_coeff is not None: + self.expert_bias_E = torch.zeros(self.num_experts, dtype=torch.float32) + + +def _mlp1_bias_init(stock_init: Callable) -> Callable: + """Stock GPT-OSS gate/up bias init, plus GPT-OSS's +1 on the up bias. + + The NKI kernel packs the gate/up bias as ``(E, 2, I)``; index 1 is up. + """ + + def init(tensor) -> None: + stock_init(tensor) + tensor[:, 1, :] += 1.0 + + return init + + +def to_nki_moe_config(moe: GptOssMoE.Config) -> NKIMoE.Config: + """Convert a stock GPT-OSS MoE config into the Neuron NKI implementation. + + Keeps the stock router gate and per-parameter initialization, so a model + built from either config starts from the same distribution. + """ + if not isinstance(moe, GptOssMoE.Config): + raise TypeError(f"Expected GptOssMoE.Config, got {type(moe)!r}") + if not isinstance(moe.router, TokenChoiceTopKRouter.Config): + raise TypeError( + f"Expected TokenChoiceTopKRouter.Config, got {type(moe.router)!r}" + ) + if not isinstance(moe.experts, GptOssGroupedExperts.Config): + raise TypeError( + f"Expected GptOssGroupedExperts.Config, got {type(moe.experts)!r}" + ) + stock_init = moe.experts.param_init + return NKIMoE.Config( + num_experts=moe.num_experts, + load_balance_coeff=moe.load_balance_coeff, + router=MaskedScoresTopKRouter.Config( + gate=replace(moe.router.gate), + top_k=moe.router.top_k, + score_func=moe.router.score_func, + route_norm=moe.router.route_norm, + ), + experts=NKIGroupedExperts.Config( + dim=moe.experts.dim, + hidden_dim=moe.experts.hidden_dim, + num_experts=moe.num_experts, + top_k=moe.router.top_k, + block_size=256, + swiglu_limit=moe.experts.swiglu_limit, + param_init={ + "mlp1_weight": stock_init["mlp1_weight_EGD"], + "mlp1_bias": _mlp1_bias_init(stock_init["mlp1_bias_EG"]), + "mlp2_weight": stock_init["mlp2_weight_EDF"], + "mlp2_bias": stock_init["mlp2_bias_ED"], + }, + ), + ) diff --git a/torchtitan_neuron/models/llama3/__init__.py b/torchtitan_neuron/models/llama3/__init__.py new file mode 100644 index 00000000000..8c1cdfdf061 --- /dev/null +++ b/torchtitan_neuron/models/llama3/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron-specific Llama3 recipes.""" diff --git a/torchtitan_neuron/models/llama3/config_registry.py b/torchtitan_neuron/models/llama3/config_registry.py new file mode 100644 index 00000000000..1aaeb74865f --- /dev/null +++ b/torchtitan_neuron/models/llama3/config_registry.py @@ -0,0 +1,97 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron Llama3 dense trainer configs.""" + +from torchtitan.components.metrics import MetricsProcessor +from torchtitan.config import CompileConfig +from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC +from torchtitan.models.llama3 import model_registry +from torchtitan.models.llama3.config_registry import llama3_70b, llama3_8b +from torchtitan.trainer import Trainer + + +def _llama3_neuron_dense_config( + flavor, base_factory, *, full_ac: bool, compile_config: CompileConfig | None = None +) -> Trainer.Config: + config = base_factory() + config.model_spec = model_registry(flavor, attn_backend="sdpa") + # Neuron does not support the RoPE gather; use the prefix-cache positions path. + # Applied here (alongside the sdpa attn backend) rather than via a dedicated + # model flavor. + for layer in config.model_spec.model.layers: + layer.attention.rope.use_prefix_cache_for_positions = True + config.metrics = MetricsProcessor.Config(log_freq=1) + config.dataloader.dataset = "c4_test" + config.activation_checkpoint = FullAC.Config() if full_ac else SelectiveAC.Config() + if compile_config is not None: + config.compile = compile_config + return config + + +def llama3_8b_sac_eager() -> Trainer.Config: + return _llama3_neuron_dense_config("8B", llama3_8b, full_ac=False) + + +def llama3_8b_fullac_eager() -> Trainer.Config: + return _llama3_neuron_dense_config("8B", llama3_8b, full_ac=True) + + +def llama3_70b_fullac_eager() -> Trainer.Config: + return _llama3_neuron_dense_config("70B", llama3_70b, full_ac=True) + + +def llama3_8b_sac_compile() -> Trainer.Config: + return _llama3_neuron_dense_config( + "8B", + llama3_8b, + full_ac=False, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + + +def llama3_8b_fullac_compile() -> Trainer.Config: + return _llama3_neuron_dense_config( + "8B", + llama3_8b, + full_ac=True, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + + +def llama3_70b_fullac_compile() -> Trainer.Config: + return _llama3_neuron_dense_config( + "70B", + llama3_70b, + full_ac=True, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + + +def llama3_70b_16layers_fullac_compile() -> Trainer.Config: + config = _llama3_neuron_dense_config( + "70B", + llama3_70b, + full_ac=True, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + config.model_spec.model.layers = config.model_spec.model.layers[:16] + return config diff --git a/torchtitan_neuron/models/qwen3/__init__.py b/torchtitan_neuron/models/qwen3/__init__.py new file mode 100644 index 00000000000..7246515b3fb --- /dev/null +++ b/torchtitan_neuron/models/qwen3/__init__.py @@ -0,0 +1,99 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron-specific Qwen3 model implementations.""" + +from dataclasses import replace + +from torchtitan.models.common import CosSinRoPE, Embedding, Linear +from torchtitan.models.qwen3 import ( + _build_qwen3_moe_layers, + _EMBEDDING_INIT, + _output_linear_init, + _qwen3_norm, + model_registry as core_model_registry, + Qwen3Model, +) +from torchtitan.models.utils import validate_converter_order +from torchtitan.protocols.model import ModelConfigConverter +from torchtitan.protocols.model_spec import ModelSpec + + +__all__ = ["model_registry"] + + +def _235b_a22b_scaled( + attn_backend: str, + moe_comm_backend: str = "standard", +) -> Qwen3Model.Config: + """Build the reduced-shape 235B validation flavor.""" + dim = 1024 + head_dim = 128 + vocab_size = 151936 + return Qwen3Model.Config( + vocab_size=vocab_size, + dim=dim, + norm=_qwen3_norm(dim), + tok_embeddings=Embedding.Config( + num_embeddings=vocab_size, + embedding_dim=dim, + param_init=_EMBEDDING_INIT, + ), + lm_head=Linear.Config( + in_features=dim, + out_features=vocab_size, + param_init=_output_linear_init(dim), + ), + layers=_build_qwen3_moe_layers( + fuse_qkv=True, + n_layers=2, + dim=dim, + n_heads=8, + n_kv_heads=2, + head_dim=head_dim, + moe_hidden_dim=512, + num_experts=16, + top_k=4, + attn_backend=attn_backend, + rope=CosSinRoPE.Config( + dim=head_dim, + max_seq_len=40960, + theta=5000000.0, + ), + moe_comm_backend=moe_comm_backend, + ), + ) + + +def model_registry( + flavor: str, + attn_backend: str = "flex", + moe_comm_backend: str | None = None, + converters: list[ModelConfigConverter.Config] | None = None, +) -> ModelSpec: + """Resolve core Qwen3 flavors plus Neuron-owned validation flavors.""" + if flavor != "235B-A22B-scaled": + return core_model_registry( + flavor, + attn_backend=attn_backend, + moe_comm_backend=moe_comm_backend, + converters=converters, + ) + + config = _235b_a22b_scaled( + attn_backend=attn_backend, + moe_comm_backend=moe_comm_backend or "standard", + ) + if converters is not None: + validate_converter_order(converters) + for converter in converters: + config = converter.build().convert(config) + + base_spec = core_model_registry( + "235B-A22B", + attn_backend=attn_backend, + moe_comm_backend=moe_comm_backend, + ) + return replace(base_spec, flavor=flavor, model=config) diff --git a/torchtitan_neuron/models/qwen3/_mxfp8_flavors.py b/torchtitan_neuron/models/qwen3/_mxfp8_flavors.py new file mode 100644 index 00000000000..80121cace83 --- /dev/null +++ b/torchtitan_neuron/models/qwen3/_mxfp8_flavors.py @@ -0,0 +1,127 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron MXFP8 Qwen3 model flavors for prequantization validation. + +These alignment-specific flavors are kept out of the core Qwen3 registry. The +recipes that use them call ``register_mxfp8_flavors()`` to register them into +core's ``qwen3_configs`` on demand, so merely importing the Neuron package does +not mutate core's registry. See the per-flavor docstrings for the alignment +arithmetic. +""" + +from functools import partial + +from torchtitan.models.common import CosSinRoPE, Embedding, Linear +from torchtitan.models.qwen3 import ( + _build_qwen3_layers, + _EMBEDDING_INIT, + _output_linear_init, + _qwen3_norm, + qwen3_configs, + Qwen3Model, +) + + +def _32b_dim6144_ffn24576(attn_backend: str, n_layers: int = 6) -> Qwen3Model.Config: + """Qwen3-32B reshaped to dim=6144 / hidden_dim=24576 (5.26B params at 6 layers). + + Every per-rank weight shard is a multiple of 128 in both dims under TP=4 x + FSDP=16 with FSDP on dim 1, meeting the MXFP8 quantize kernel's 128-element + floor on the blocked dim; requires SimpleFSDP dim-1 sharding support. Alignment + is per-weight, so n_layers does not affect it (6 layers keeps seed-checkpoint + numerics runs cheap). Not a real Qwen3-32B shape and not a convergence target; + lm_head and tok_embeddings keep the stock vocab_size, outside the alignment + claim (the MXFP8 converters scope to fqns=["layers"]). + """ + dim = 6144 + head_dim = 128 + vocab_size = 151936 + return Qwen3Model.Config( + vocab_size=vocab_size, + dim=dim, + norm=_qwen3_norm(dim), + tok_embeddings=Embedding.Config( + num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT + ), + lm_head=Linear.Config( + in_features=dim, + out_features=vocab_size, + param_init=_output_linear_init(dim), + ), + layers=_build_qwen3_layers( + fuse_qkv=True, + n_layers=n_layers, + dim=dim, + n_heads=64, + n_kv_heads=8, + head_dim=head_dim, + hidden_dim=24576, + attn_backend=attn_backend, + rope=CosSinRoPE.Config( + dim=head_dim, + max_seq_len=40960, + theta=1000000.0, + ), + ), + ) + + +def _32b_dim8192_ffn32768(attn_backend: str, n_layers: int = 6) -> Qwen3Model.Config: + """Qwen3-32B reshaped so every per-rank K is 512 (10.94B params at 6 layers). + + The wrapX packing needs each rank's local K to be a whole multiple of the quantize + kernel's 512-element tile, otherwise the per-rank encodings do not concatenate into + the layout whole-K quantization produces and the matmul silently contracts mismatched + K (see ``MIN_K`` in ``mx_prequant.py``). These shapes satisfy that for all weights + under TP=4 x FSDP=16 on dim 1; requires SimpleFSDP dim-1 sharding support. Not a real + Qwen3-32B shape and not a convergence target; lm_head and tok_embeddings keep the stock + vocab_size, outside the alignment claim (the MXFP8 converters scope to fqns=["layers"]). + """ + dim = 8192 + head_dim = 128 + vocab_size = 151936 + return Qwen3Model.Config( + vocab_size=vocab_size, + dim=dim, + norm=_qwen3_norm(dim), + tok_embeddings=Embedding.Config( + num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT + ), + lm_head=Linear.Config( + in_features=dim, + out_features=vocab_size, + param_init=_output_linear_init(dim), + ), + layers=_build_qwen3_layers( + fuse_qkv=True, + n_layers=n_layers, + dim=dim, + n_heads=256, + n_kv_heads=32, + head_dim=head_dim, + hidden_dim=32768, + attn_backend=attn_backend, + rope=CosSinRoPE.Config( + dim=head_dim, + max_seq_len=40960, + theta=1000000.0, + ), + ), + ) + + +# Alignment-validation flavors, keyed as model_registry(flavor) expects. Registered +# on demand by the recipes that use them (register_mxfp8_flavors), not at import. +_MXFP8_FLAVORS = { + "32B_dim6144_ffn24576_6layers": partial(_32b_dim6144_ffn24576, n_layers=6), + "32B_dim8192_ffn32768_6layers": partial(_32b_dim8192_ffn32768, n_layers=6), +} + + +def register_mxfp8_flavors() -> None: + """Register the MXFP8 validation flavors in core's qwen3_configs.""" + for name, builder in _MXFP8_FLAVORS.items(): + qwen3_configs.setdefault(name, builder) diff --git a/torchtitan_neuron/models/qwen3/config_registry.py b/torchtitan_neuron/models/qwen3/config_registry.py new file mode 100644 index 00000000000..5c2d2c9f0d7 --- /dev/null +++ b/torchtitan_neuron/models/qwen3/config_registry.py @@ -0,0 +1,670 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron Qwen3 dense and NKI MoE trainer configs.""" + +from collections.abc import Callable +from dataclasses import dataclass + +from torchtitan.components.checkpoint import CheckpointManager +from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss +from torchtitan.components.lr_scheduler import LRSchedulersContainer +from torchtitan.components.metrics import MetricsProcessor +from torchtitan.components.optimizer import default_adamw +from torchtitan.components.quantization import MXFP8LinearConverter +from torchtitan.config import CompileConfig, derive, ParallelismConfig, TrainingConfig +from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC +from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader +from torchtitan.models.common.config_utils import decoder_vocab_size +from torchtitan.models.qwen3.config_registry import qwen3_30b_a3b, qwen3_32b +from torchtitan.protocols.model_spec import ModelSpec +from torchtitan.tools.profiler import Profiler +from torchtitan.trainer import Trainer + +from . import model_registry +from .nki_moe import BWMMGroupedExperts, to_nki_moe_config + + +def qwen3_235b_a22b() -> Trainer.Config: + model_spec = model_registry("235B-A22B") + return Trainer.Config( + loss=ChunkedLossWrapper.Config( + loss_fn=CrossEntropyLoss.Config( + global_vocab_size=decoder_vocab_size(model_spec), + ), + ), + hf_assets_path="./assets/hf/Qwen3-235B-A22B", + model_spec=model_spec, + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"), + optimizer=default_adamw(lr=8e-4), + lr_scheduler=LRSchedulersContainer.Config(warmup_steps=600), + training=TrainingConfig( + local_batch_size=2, + seq_len=4096, + steps=3000, + ), + parallelism=ParallelismConfig( + data_parallel_shard_degree=-1, + tensor_parallel_degree=1, + context_parallel_degree=1, + pipeline_parallel_degree=1, + ), + checkpoint=CheckpointManager.Config( + interval=500, + last_save_model_only=False, + export_dtype="float16", + ), + activation_checkpoint=FullAC.Config(), + ) + + +@dataclass(frozen=True) +class Qwen3NKIMoEPreset: + config_factory: Callable[[], Trainer.Config] + base_flavor: str + flavor: str + + +_QWEN3_NKI_MOE_PRESETS = { + "30B-A3B": Qwen3NKIMoEPreset( + config_factory=qwen3_30b_a3b, + base_flavor="30B-A3B", + flavor="30B-A3B-nki", + ), + "235B-A22B": Qwen3NKIMoEPreset( + config_factory=qwen3_235b_a22b, + base_flavor="235B-A22B", + flavor="235B-A22B-nki", + ), +} + + +def _reduce_moe_experts(config, *, num_experts: int, top_k: int) -> None: + for layer in config.model_spec.model.layers: + moe = getattr(layer, "moe", None) + if moe is None: + continue + moe.num_experts = num_experts + moe.router.num_experts = num_experts + moe.router.top_k = top_k + moe.router.gate.out_features = num_experts + moe.experts.num_experts = num_experts + moe.experts.token_dispatcher.num_experts = num_experts + moe.experts.token_dispatcher.top_k = top_k + + +def select_neuron_components(model_spec: ModelSpec) -> None: + for layer in model_spec.model.layers: + layer.attention.rope.use_prefix_cache_for_positions = True + moe = getattr(layer, "moe", None) + if moe is not None: + layer.moe = to_nki_moe_config(moe) + + +def qwen3_nki_moe_config( + preset_name: str, + *, + compile_config: CompileConfig | None, + n_layers: int | None = None, + num_experts: int | None = None, + top_k: int | None = None, +) -> Trainer.Config: + if (num_experts is None) != (top_k is None): + raise ValueError("num_experts and top_k must be set together") + if num_experts is not None and top_k is not None and top_k > num_experts: + raise ValueError(f"top_k={top_k} exceeds num_experts={num_experts}") + + preset = _QWEN3_NKI_MOE_PRESETS[preset_name] + config = preset.config_factory() + config.model_spec = model_registry(preset.base_flavor, attn_backend="sdpa") + + flavor = preset.flavor + if num_experts is not None: + flavor = f"{flavor}-e{num_experts}-topk{top_k}" + if n_layers is not None: + flavor = f"{flavor}-{n_layers}l" + config.model_spec.model.layers = config.model_spec.model.layers[:n_layers] + config.model_spec.flavor = flavor + + if num_experts is not None: + assert top_k is not None + _reduce_moe_experts(config, num_experts=num_experts, top_k=top_k) + select_neuron_components(config.model_spec) + + config.compile = compile_config or CompileConfig(enable=False, backend="neuron") + config.override.imports = [] + config.training.steps = 10 + config.dataloader.dataset = "c4_test" + config.lr_scheduler.warmup_steps = 10 + for group in config.optimizer.param_groups: + group.optimizer_kwargs["lr"] = 3e-4 + config.metrics.log_freq = 1 + return config + + +def qwen3_30b_a3b_nki_fullac_eager() -> Trainer.Config: + return qwen3_nki_moe_config( + "30B-A3B", + compile_config=CompileConfig(enable=False, backend="neuron"), + ) + + +def qwen3_30b_a3b_nki_fullac_compile() -> Trainer.Config: + return qwen3_nki_moe_config( + "30B-A3B", + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + + +def qwen3_235b_a22b_neuron() -> Trainer.Config: + """Full-scale Qwen3-235B Neuron recipe.""" + config = qwen3_235b_a22b() + config.model_spec = model_registry( + "235B-A22B", + attn_backend="sdpa", + moe_comm_backend="neuron_ep", + ) + config.model_spec.flavor = "235B-A22B-neuron" + for layer in config.model_spec.model.layers: + layer.moe.experts = derive( + layer.moe.experts, + BWMMGroupedExperts.Config, + ) + + config.training.local_batch_size = 1 + config.training.seq_len = 4096 + config.parallelism = ParallelismConfig( + data_parallel_replicate_degree=4, + data_parallel_shard_degree=32, + tensor_parallel_degree=4, + pipeline_parallel_degree=2, + expert_parallel_degree=128, + ) + config.activation_checkpoint = SelectiveAC.Config() + config.compile = CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ) + config.override.imports = [] + return config + + +def qwen3_8b_tp4_fsdp16_sac() -> Trainer.Config: + model_spec = model_registry("8B", attn_backend="sdpa") + config = Trainer.Config( + loss=ChunkedLossWrapper.Config( + loss_fn=CrossEntropyLoss.Config( + global_vocab_size=decoder_vocab_size(model_spec), + ), + ), + hf_assets_path="./assets/hf/Qwen3-8B", + model_spec=model_spec, + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"), + optimizer=default_adamw(lr=8e-4), + lr_scheduler=LRSchedulersContainer.Config(warmup_steps=10), + training=TrainingConfig(local_batch_size=1, seq_len=4096, steps=50), + parallelism=ParallelismConfig( + data_parallel_shard_degree=-1, + tensor_parallel_degree=4, + context_parallel_degree=1, + pipeline_parallel_degree=1, + ), + checkpoint=CheckpointManager.Config( + interval=500, + last_save_model_only=False, + export_dtype="float16", + ), + activation_checkpoint=SelectiveAC.Config(), + ) + for layer in config.model_spec.model.layers: + layer.attention.rope.use_prefix_cache_for_positions = True + return config + + +def qwen3_8b_tp4_fsdp16_compile_sac() -> Trainer.Config: + config = qwen3_8b_tp4_fsdp16_sac() + config.compile = CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ) + return config + + +def _qwen3_neuron_dense_config( + flavor: str, + *, + tensor_parallel_degree: int, + global_batch_size: int, + compile_config: CompileConfig | None, + selective_ac: bool = False, + lr: float | None = None, +) -> Trainer.Config: + config = qwen3_32b() + # The base is qwen3_32b() (lr 8e-4). 0.6B uses a different lr (3e-4), so allow + # a per-size override. + if lr is not None: + config.optimizer = default_adamw(lr=lr) + config.model_spec = model_registry(flavor, attn_backend="sdpa") + # Preserve the sequential-position behavior selected by these recipes. + for layer in config.model_spec.model.layers: + layer.attention.rope.use_prefix_cache_for_positions = True + config.hf_assets_path = f"./assets/hf/Qwen3-{flavor}" + config.metrics = MetricsProcessor.Config(log_freq=1) + config.dataloader.dataset = "c4_test" + config.training.local_batch_size = 1 + config.training.steps = 10 + config.training.global_batch_size = global_batch_size + config.parallelism.tensor_parallel_degree = tensor_parallel_degree + config.parallelism.data_parallel_shard_degree = -1 + config.parallelism.context_parallel_degree = 1 + config.parallelism.pipeline_parallel_degree = 1 + config.activation_checkpoint = ( + SelectiveAC.Config() if selective_ac else FullAC.Config() + ) + config.compile = compile_config or CompileConfig() + return config + + +def qwen3_0_6b_tp1_fsdp64_sac_eager() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "0.6B", + tensor_parallel_degree=1, + global_batch_size=64, + compile_config=None, + selective_ac=True, + lr=3e-4, + ) + + +def qwen3_0_6b_tp1_fsdp64_fullac_compile() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "0.6B", + tensor_parallel_degree=1, + global_batch_size=64, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + lr=3e-4, + ) + + +def qwen3_1_7b_tp1_fsdp64_fullac_eager() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "1.7B", tensor_parallel_degree=1, global_batch_size=64, compile_config=None + ) + + +def qwen3_1_7b_tp1_fsdp64_fullac_compile() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "1.7B", + tensor_parallel_degree=1, + global_batch_size=64, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + + +def qwen3_4b_tp4_fsdp16_sac_eager() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "4B", + tensor_parallel_degree=4, + global_batch_size=16, + compile_config=None, + selective_ac=True, + ) + + +def qwen3_4b_tp4_fsdp16_sac_compile() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "4B", + tensor_parallel_degree=4, + global_batch_size=16, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + selective_ac=True, + ) + + +def qwen3_8b_tp1_fsdp64_fullac_eager() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "8B", tensor_parallel_degree=1, global_batch_size=64, compile_config=None + ) + + +def qwen3_8b_tp1_fsdp64_fullac_compile() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "8B", + tensor_parallel_degree=1, + global_batch_size=64, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + + +def qwen3_14b_tp4_fsdp16_fullac_eager() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "14B", tensor_parallel_degree=4, global_batch_size=16, compile_config=None + ) + + +def qwen3_14b_tp4_fsdp16_fullac_compile() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "14B", + tensor_parallel_degree=4, + global_batch_size=16, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + + +def qwen3_32b_tp4_fsdp16_fullac_eager() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "32B", tensor_parallel_degree=4, global_batch_size=16, compile_config=None + ) + + +def qwen3_32b_tp4_fsdp16_fullac_compile() -> Trainer.Config: + return _qwen3_neuron_dense_config( + "32B", + tensor_parallel_degree=4, + global_batch_size=16, + compile_config=CompileConfig( + enable=True, + components=["model", "loss", "grad_clip"], + backend="neuron", + ), + ) + + +def _qwen3_235b_4layers( + moe_comm_backend: str, + n_layers: int = 4, + ep_degree: int = 4, + num_experts: int | None = None, + top_k: int | None = None, + tp_degree: int = 1, +) -> Trainer.Config: + """Build a reduced-layer Qwen3-235B expert-parallel recipe.""" + if (num_experts is None) != (top_k is None): + raise ValueError("num_experts and top_k must be set together.") + model_spec = model_registry( + "235B-A22B", + attn_backend="sdpa", + moe_comm_backend=moe_comm_backend, + ) + model_spec.model.layers = model_spec.model.layers[:n_layers] + if num_experts is not None: + assert top_k is not None + for layer in model_spec.model.layers: + moe = layer.moe + moe.num_experts = num_experts + moe.router.num_experts = num_experts + moe.router.top_k = top_k + moe.router.gate.out_features = num_experts + moe.experts.num_experts = num_experts + moe.experts.token_dispatcher.num_experts = num_experts + moe.experts.token_dispatcher.top_k = top_k + if moe_comm_backend == "neuron_ep": + for layer in model_spec.model.layers: + layer.moe.experts = derive( + layer.moe.experts, + BWMMGroupedExperts.Config, + ) + return Trainer.Config( + loss=ChunkedLossWrapper.Config( + loss_fn=CrossEntropyLoss.Config( + global_vocab_size=decoder_vocab_size(model_spec), + ), + ), + hf_assets_path="./tests/assets/tokenizer", + metrics=MetricsProcessor.Config(log_freq=1), + model_spec=model_spec, + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4_test"), + optimizer=default_adamw(lr=3e-4), + lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), + training=TrainingConfig(local_batch_size=1, seq_len=512, steps=5), + parallelism=ParallelismConfig( + expert_parallel_degree=ep_degree, + tensor_parallel_degree=tp_degree, + ), + checkpoint=CheckpointManager.Config( + interval=1000, + last_save_model_only=False, + export_dtype="float16", + ), + activation_checkpoint=None, + ) + + +def _scaled_235b_config(moe_comm_backend: str, ep_degree: int) -> Trainer.Config: + model_spec = model_registry( + "235B-A22B-scaled", + attn_backend="sdpa", + moe_comm_backend=moe_comm_backend, + ) + return Trainer.Config( + loss=ChunkedLossWrapper.Config( + loss_fn=CrossEntropyLoss.Config( + global_vocab_size=decoder_vocab_size(model_spec), + ), + ), + hf_assets_path="./tests/assets/tokenizer", + metrics=MetricsProcessor.Config(log_freq=1), + model_spec=model_spec, + dataloader=HuggingFaceTextDataLoader.Config(dataset="c4_test"), + optimizer=default_adamw(lr=3e-4), + lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), + training=TrainingConfig(local_batch_size=1, seq_len=512, steps=5), + parallelism=ParallelismConfig(expert_parallel_degree=ep_degree), + checkpoint=CheckpointManager.Config( + interval=1000, + last_save_model_only=False, + export_dtype="float16", + ), + activation_checkpoint=None, + ) + + +def qwen3_235b_scaled_singlerank() -> Trainer.Config: + return _scaled_235b_config("neuron_ep", ep_degree=1) + + +def _qwen3_235b_scaled_ep4(moe_comm_backend: str) -> Trainer.Config: + return _scaled_235b_config(moe_comm_backend, ep_degree=4) + + +def qwen3_235b_scaled_ep4_standard() -> Trainer.Config: + return _qwen3_235b_scaled_ep4(moe_comm_backend="standard") + + +def qwen3_235b_scaled_64rank_noep() -> Trainer.Config: + return _scaled_235b_config("standard", ep_degree=1) + + +def qwen3_235b_4layers_neuron_standard() -> Trainer.Config: + return _qwen3_235b_4layers(moe_comm_backend="standard") + + +def qwen3_235b_2layers_neuron_standard() -> Trainer.Config: + return _qwen3_235b_4layers(moe_comm_backend="standard", n_layers=2) + + +def qwen3_235b_4layers_ep64_nle1_neuron_ep() -> Trainer.Config: + return _qwen3_235b_4layers( + moe_comm_backend="neuron_ep", + ep_degree=64, + num_experts=64, + top_k=8, + ) + + +def qwen3_235b_4layers_ep64_nle1_neuron_ep_profile() -> Trainer.Config: + config = qwen3_235b_4layers_ep64_nle1_neuron_ep() + config.training.steps = 8 + config.profiler = Profiler.Config( + enable_profiling=True, + profile_freq=6, + profiler_warmup=2, + profiler_active=1, + ) + return config + + +def qwen3_235b_4layers_ep64_nle1_tp4_neuron_ep() -> Trainer.Config: + return _qwen3_235b_4layers( + moe_comm_backend="neuron_ep", + ep_degree=64, + num_experts=64, + top_k=8, + tp_degree=4, + ) + + +def qwen3_235b_4layers_ep64_nle1_tp4_neuron_ep_compile() -> Trainer.Config: + config = qwen3_235b_4layers_ep64_nle1_tp4_neuron_ep() + config.compile = CompileConfig( + enable=True, + components=["loss"], + backend="neuron", + ) + return config + + +def qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep() -> Trainer.Config: + config = _qwen3_235b_4layers( + moe_comm_backend="neuron_ep", + n_layers=6, + ep_degree=64, + num_experts=64, + top_k=8, + tp_degree=4, + ) + config.training.seq_len = 4096 + config.activation_checkpoint = SelectiveAC.Config() + return config + + +def qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_compile() -> Trainer.Config: + config = qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep() + config.compile = CompileConfig( + enable=True, + components=["model", "loss"], + backend="neuron", + ) + return config + + +def qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_compile() -> Trainer.Config: + config = _qwen3_235b_4layers( + moe_comm_backend="neuron_ep", + n_layers=6, + ep_degree=128, + num_experts=128, + top_k=8, + tp_degree=4, + ) + config.training.seq_len = 4096 + config.activation_checkpoint = SelectiveAC.Config() + config.compile = CompileConfig( + enable=True, + components=["model", "loss"], + backend="neuron", + ) + return config + + +def qwen3_235b_94layers_ep128_tp4_fsdp64_pp1_4k_neuron_ep_compile() -> (Trainer.Config): + config = _qwen3_235b_4layers( + moe_comm_backend="neuron_ep", + n_layers=94, + ep_degree=128, + num_experts=128, + top_k=8, + tp_degree=4, + ) + config.training.seq_len = 4096 + config.training.local_batch_size = 1 + config.activation_checkpoint = FullAC.Config() + config.compile = CompileConfig( + enable=True, + components=["model", "loss"], + backend="neuron", + ) + config.parallelism.data_parallel_replicate_degree = 1 + config.parallelism.data_parallel_shard_degree = 64 + config.parallelism.pipeline_parallel_degree = 1 + return config + + +def qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_pp2_compile() -> Trainer.Config: + config = qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_compile() + config.parallelism.pipeline_parallel_degree = 2 + config.parallelism.pipeline_parallel_schedule = "1F1B" + config.parallelism.pipeline_parallel_first_stage_less_layers = 0 + config.parallelism.pipeline_parallel_last_stage_less_layers = 0 + config.training.local_batch_size = 16 + return config + + +def qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_pp2_compile() -> Trainer.Config: + config = qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_compile() + config.parallelism.pipeline_parallel_degree = 2 + config.parallelism.pipeline_parallel_schedule = "1F1B" + config.parallelism.pipeline_parallel_first_stage_less_layers = 0 + config.parallelism.pipeline_parallel_last_stage_less_layers = 0 + config.training.local_batch_size = 16 + return config + + +def qwen3_235b_4layers_ep64_neuron_standard() -> Trainer.Config: + return _qwen3_235b_4layers(moe_comm_backend="standard", ep_degree=64) + + +# --- Neuron trn3 dense qwen3-8B MXFP8 (eager by default) --- +# Gating, backend selection, and the torch.compile path: see +# torchtitan_neuron/docs/neuron_mxfp8.md. + + +def qwen3_8b_mxfp8() -> Trainer.Config: + """qwen3-8B MXFP8 on all in-layer linears (attention + FFN); LM head stays bf16. + + Inherits ``qwen3_8b_tp4_fsdp16_sac`` (SDPA + prefix RoPE cache, TP4 x FSDP16, + SelectiveAC) and only attaches ``MXFP8LinearConverter`` (``fqns=["layers"]``, + ``kernel_preference="neuron"``). + """ + config = qwen3_8b_tp4_fsdp16_sac() + model_compile_enabled = ( + config.compile.enable and "model" in config.compile.components + ) + config.model_spec.model = ( + MXFP8LinearConverter.Config( + fqns=["layers"], + kernel_preference="neuron", + model_compile_enabled=model_compile_enabled, + ) + .build() + .convert(config.model_spec.model) + ) + return config diff --git a/torchtitan_neuron/models/qwen3/nki_moe.py b/torchtitan_neuron/models/qwen3/nki_moe.py new file mode 100644 index 00000000000..61eac11a06a --- /dev/null +++ b/torchtitan_neuron/models/qwen3/nki_moe.py @@ -0,0 +1,795 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Qwen3 NKI MoE implementation for the Neuron extension package.""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Literal + +import nki.language as nl + +import spmd_types as spmd +import torch +import torch.nn.functional as F +from nkilib.core.subkernels.find_nonzero_indices import find_nonzero_indices +from nkilib.core.subkernels.indexed_flatten import indexed_flatten +from nkilib.core.utils.common_types import ActFnType, ExpertAffinityScaleMode +from nkilib.experimental.moe.bwd.blockwise_mm_backward import blockwise_mm_bwd +from nkilib.experimental.moe.bwd.moe_bwd_parameters import ( + AffinityOption, + ShardOption, + SkipMode as BwdSkipMode, +) +from nkilib.experimental.moe.forward.bwmm_shard_on_H import ( + blockwise_mm_baseline_shard_hidden, + SkipMode as FwdSkipMode, +) +from torch import nn +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Partial, Replicate +from torch_neuronx import nki_op +from torch_neuronx.nki_hop import wrap_nki +from torch_neuronx.utils import get_logical_neuron_cores + +from torchtitan.models.common import Linear +from torchtitan.models.common.decoder_sharding import ( + dense_activation_placement, + dense_param_placement, + dense_sequence_parallel_placement, +) +from torchtitan.models.common.moe import ( + attach_global_batch_lbl, + GroupedExperts, + MoE, + TokenChoiceTopKRouter, +) +from torchtitan.protocols.module import Module +from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig + + +BLOCK_SIZE = 512 +EP_BLOCK_SIZE = 256 +LNC = (int(get_logical_neuron_cores()),) + +wrapped_moe_fwd = wrap_nki(blockwise_mm_baseline_shard_hidden) +wrapped_moe_bwd = wrap_nki(blockwise_mm_bwd) +wrapped_find_nonzero = wrap_nki(find_nonzero_indices) +wrapped_indexed_flatten = wrap_nki(indexed_flatten) + + +def _single_expert_block_metadata( + tokens: int, + num_blocks: int, + block_size: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + token_position_to_id = torch.arange( + num_blocks * block_size, + device=device, + dtype=torch.int32, + ) + token_position_to_id = torch.where( + token_position_to_id < tokens, + token_position_to_id, + -1, + ) + block_to_expert = torch.zeros( + (num_blocks, 1), + device=device, + dtype=torch.int32, + ) + return token_position_to_id, block_to_expert + + +@nki_op("qwen3_nki_moe::nki_moe_fwd") +def nki_fwd( + hidden_states: torch.Tensor, + masked_scores: torch.Tensor, + gate_up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + top_k: int, + block_size: int, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + int, + int, + torch.Tensor, +]: + tokens = hidden_states.shape[0] + num_experts = masked_scores.shape[1] + hidden_dim_tp = down_proj_weight.shape[1] + # At most one affinity can exist per local expert for a token row. For + # neuron_ep, masked_scores describes the post-dispatch local experts. + max_affinities_per_row = min(top_k, num_experts) + num_blocks = ( + tokens * max_affinities_per_row + block_size - 1 + ) // block_size + num_experts + + if num_experts == 1: + # find_nonzero_indices requires the expert dimension to be divisible + # across LNC2. Dispatched nle=1 inputs are already expert-major, so the + # block metadata is an identity mapping and needs no NKI preprocessing. + token_position_to_id, block_to_expert = _single_expert_block_metadata( + tokens, + num_blocks, + block_size, + hidden_states.device, + ) + else: + indices, counts = wrapped_find_nonzero[LNC](masked_scores.float()) + blocks_per_expert = ((counts + block_size - 1) // block_size).clamp(min=1) + slack = num_blocks - blocks_per_expert.sum() + padded_blocks_per_expert = blocks_per_expert.clone() + padded_blocks_per_expert[-1] += slack + row_offsets = ( + torch.cumsum(padded_blocks_per_expert, dim=0) - padded_blocks_per_expert + ) + expert_ids = torch.arange( + num_experts, device=hidden_states.device, dtype=torch.int32 + ) + # Repeat expert ids by padded block count. Repeating the counts themselves + # can create invalid expert ids and drive the NKI gather out of bounds. + block_to_expert = torch.repeat_interleave( + expert_ids, padded_blocks_per_expert, output_size=num_blocks + ).unsqueeze(1) + token_position_to_id = wrapped_indexed_flatten[LNC]( + indices, + block_size, + num_blocks * block_size, + row_offsets, + padding_val=-1, + ) + + affinities = masked_scores.reshape(-1, 1).clone() + gate_up_activations = torch.empty( + (num_blocks, 2, hidden_dim_tp, block_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + down_activations = None + output, gate_up_activations = wrapped_moe_fwd[LNC]( + hidden_states, + affinities, + gate_up_proj_weight, + down_proj_weight, + block_size, + token_position_to_id, + block_to_expert, + gate_up_activations, + down_activations, + skip_dma=FwdSkipMode(skip_token=True), + is_tensor_update_accumulating=True, + expert_affinities_scaling_mode=ExpertAffinityScaleMode.POST_SCALE, + # The kernel sums the top-k expert outputs in bf16 unless told otherwise. + accum_dtype=nl.float32, + ) + + # Inputs are recovered from setup_context. Returning them would violate + # the custom-op no-alias contract enforced by PT2.12. + return ( + affinities, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + tokens, + num_experts, + output, + ) + + +@nki_op( + "qwen3_nki_moe::nki_moe_bwd", + mutates_args={ + "hidden_states_grad_out", + "expert_affinities_grad_out", + "gate_up_proj_weight_grad_out", + "down_proj_weight_grad_out", + }, +) +def nki_bwd( + hidden_states: torch.Tensor, + affinities: torch.Tensor, + gate_up_proj_weight: torch.Tensor, + down_proj_weight: torch.Tensor, + gate_up_activations: torch.Tensor, + down_activations: torch.Tensor, + token_position_to_id: torch.Tensor, + block_to_expert: torch.Tensor, + grad_output: torch.Tensor, + hidden_states_grad_out: torch.Tensor, + expert_affinities_grad_out: torch.Tensor, + gate_up_proj_weight_grad_out: torch.Tensor, + down_proj_weight_grad_out: torch.Tensor, + block_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return wrapped_moe_bwd[LNC]( + hidden_states, + affinities, + gate_up_proj_weight, + down_proj_weight, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + grad_output, + block_size, + skip_dma=BwdSkipMode(skip_token=True), + skip_grad_initialization=True, + activation_type=ActFnType.SiLU, + affinity_option=AffinityOption.AFFINITY_ON_I, + shard_option=ShardOption.SHARD_ON_HIDDEN, + bias=False, + hidden_states_grad_out=hidden_states_grad_out, + expert_affinities_masked_grad_out=expert_affinities_grad_out, + gate_up_proj_weight_grad_out=gate_up_proj_weight_grad_out, + down_proj_weight_grad_out=down_proj_weight_grad_out, + ) + + +def _setup_context(ctx, inputs, output): + ( + hidden_states, + masked_scores, + gate_up_proj_weight, + down_proj_weight, + _, + block_size, + ) = inputs + ( + affinities, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + tokens, + num_experts, + _, + ) = output + ctx.save_for_backward( + hidden_states, + affinities, + gate_up_proj_weight, + down_proj_weight, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + ) + ctx.tokens = tokens + ctx.num_experts = num_experts + ctx.scores_dtype = masked_scores.dtype + ctx.block_size = block_size + + +def _backward(ctx, *grad_outputs): + grad_output = grad_outputs[-1] + ( + hidden_states, + affinities, + gate_up_proj_weight, + down_proj_weight, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + ) = ctx.saved_tensors + + hidden_states_dtype = hidden_states.dtype + gate_up_proj_weight_dtype = gate_up_proj_weight.dtype + down_proj_weight_dtype = down_proj_weight.dtype + + hidden_states_grad_out = torch.zeros_like(hidden_states, dtype=torch.float32) + expert_affinities_grad_out = torch.zeros_like(affinities) + gate_up_proj_weight_grad_out = torch.zeros_like( + gate_up_proj_weight, dtype=torch.float32 + ) + down_proj_weight_grad_out = torch.zeros_like(down_proj_weight, dtype=torch.float32) + ( + hidden_states_grad, + affinities_grad, + gate_up_proj_weight_grad, + down_proj_weight_grad, + ) = nki_bwd( + hidden_states, + affinities, + gate_up_proj_weight, + down_proj_weight, + gate_up_activations, + down_activations, + token_position_to_id, + block_to_expert, + grad_output[: ctx.tokens], + hidden_states_grad_out, + expert_affinities_grad_out, + gate_up_proj_weight_grad_out, + down_proj_weight_grad_out, + ctx.block_size, + ) + hidden_states_grad = hidden_states_grad.to(hidden_states_dtype) + gate_up_proj_weight_grad = gate_up_proj_weight_grad.to(gate_up_proj_weight_dtype) + down_proj_weight_grad = down_proj_weight_grad.to(down_proj_weight_dtype) + + scores_grad = affinities_grad.reshape(-1, ctx.num_experts)[: ctx.tokens].to( + ctx.scores_dtype + ) + return ( + hidden_states_grad[: ctx.tokens], + scores_grad, + gate_up_proj_weight_grad, + down_proj_weight_grad, + None, + None, + ) + + +nki_fwd.register_autograd(_backward, setup_context=_setup_context) + + +def _debug_force_load_balance_masked_scores( + scores: torch.Tensor, + top_k: int, + num_experts: int, + route_norm: bool, + route_scale: float, +) -> torch.Tensor: + """Round-robin expert assignment with real gate scores (debug only). + + Same deterministic assignment as the GPU/std-MoE debug helper: + expert(token i, slot k) = (i * top_k + k) % num_experts. + """ + # DTensor inputs crash torch.Tensor.gather with a mixed DTensor/plain + # index. Compute on the local tensor and rewrap with the original + # (replicated) placements, which is exact for replicated scores. + _dtensor_spec = None + if isinstance(scores, DTensor): + _dtensor_spec = (scores.device_mesh, scores.placements) + scores = scores.to_local() + n_tokens = scores.shape[0] + indices = ( + torch.arange(n_tokens * top_k, device=scores.device, dtype=torch.int64).reshape( + n_tokens, top_k + ) + % num_experts + ) + top_scores = scores.gather(dim=1, index=indices) + if route_norm: + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20) + top_scores = top_scores * route_scale + + masked_scores = torch.zeros_like(scores) + masked_scores.scatter_(1, indices, top_scores) + if _dtensor_spec is not None: + masked_scores = DTensor.from_local( + masked_scores, _dtensor_spec[0], _dtensor_spec[1] + ) + return masked_scores + + +class MaskedScoresTopKRouter(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + gate: Linear.Config + top_k: int = 8 + score_func: Literal["softmax", "sigmoid"] = "softmax" + route_norm: bool = True + route_scale: float = 1.0 + _debug_force_load_balance: bool = False + + def __init__(self, config: Config): + super().__init__() + self.gate = config.gate.build() + self.top_k = config.top_k + self.score_func = config.score_func + self.route_norm = config.route_norm + self.route_scale = config.route_scale + self._debug_force_load_balance = config._debug_force_load_balance + + def forward( + self, + x: torch.Tensor, + expert_bias: torch.Tensor | None = None, + return_scores: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Masked routing scores; with ``return_scores`` also the full unbiased scores.""" + with torch.autocast(device_type=x.device.type, dtype=torch.float32): + scores = self.gate(x) + if self.score_func == "softmax": + scores = F.softmax(scores, dim=-1) + elif self.score_func == "sigmoid": + scores = torch.sigmoid(scores) + else: + raise ValueError(f"Unsupported score function: {self.score_func}") + + # Auxiliary-loss-free load balancing (https://arxiv.org/abs/2408.15664): + # the expert bias only steers expert CHOICE; gating values stay the + # unbiased scores. Same contract as std MoE / DeepSeek NKI MoE. + scores_for_choice = scores if expert_bias is None else scores + expert_bias + top_indices = torch.topk( + scores_for_choice, k=self.top_k, dim=-1, sorted=False + ).indices + mask = torch.zeros_like(scores, dtype=torch.bool) + mask.scatter_(-1, top_indices, True) + masked_scores = scores * mask + if self.route_norm: + masked_scores /= masked_scores.sum(dim=-1, keepdim=True) + 1e-20 + masked_scores = masked_scores * self.route_scale + return (masked_scores, scores) if return_scores else masked_scores + + +class NKIGroupedExperts(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + hidden_dim: int + num_experts: int + top_k: int + token_dispatcher: object | None = None + + def __init__(self, config: Config): + super().__init__() + self.num_experts = config.num_experts + self.top_k = config.top_k + self.gate_up_proj_weight = nn.Parameter( + torch.empty(config.num_experts, config.dim, 2, config.hidden_dim) + ) + self.down_proj_weight = nn.Parameter( + torch.empty(config.num_experts, config.hidden_dim, config.dim) + ) + + def forward( + self, + hidden_states: torch.Tensor, + masked_scores: torch.Tensor, + ) -> torch.Tensor: + gate_up_proj_weight = self.gate_up_proj_weight + down_proj_weight = self.down_proj_weight + if isinstance(gate_up_proj_weight, DTensor): + gate_up_proj_weight = gate_up_proj_weight.to_local() + down_proj_weight = down_proj_weight.to_local() + (*_, tokens, _, output) = nki_fwd( + hidden_states, + masked_scores, + gate_up_proj_weight, + down_proj_weight, + self.top_k, + BLOCK_SIZE, + ) + return output[:tokens] + + +class BWMMGroupedExperts(GroupedExperts): + """GroupedExperts using BWMM after the neuron_ep dispatch. + + The dispatcher returns expert-major local blocks. A one-hot local affinity + matrix routes each block to its corresponding packed expert weights; router + affinities remain applied by the dispatcher during combine. + """ + + @dataclass(kw_only=True, slots=True) + class Config(GroupedExperts.Config): + pass + + def _experts_forward( + self, + x_RD: torch.Tensor, + num_tokens_per_expert_E: torch.Tensor, + ) -> torch.Tensor: + if isinstance(self.w1_EFD, DTensor): + w1_EFD = self.w1_EFD.to_local() + assert isinstance(self.w2_EDF, DTensor) + w2_EDF = self.w2_EDF.to_local() + assert isinstance(self.w3_EFD, DTensor) + w3_EFD = self.w3_EFD.to_local() + else: + w1_EFD = self.w1_EFD + w2_EDF = self.w2_EDF + w3_EFD = self.w3_EFD + + local_experts = w1_EFD.shape[0] + if num_tokens_per_expert_E.numel() != local_experts: + raise ValueError( + "BWMMGroupedExperts requires one token count for each local " + f"expert; got {local_experts} local weight shards and " + f"{num_tokens_per_expert_E.numel()} token counts." + ) + + # BWMM stores gate/up together as [E, D, 2, H] and down as [E, H, D]. + gate_up_proj_weight = torch.stack( + (w1_EFD.transpose(-2, -1), w3_EFD.transpose(-2, -1)), + dim=2, + ).bfloat16() + down_proj_weight = w2_EDF.transpose(-2, -1).bfloat16() + + # TODO: Validate the nle>1 neuron_ep + BWMM path on hardware. + expert_ends_E = num_tokens_per_expert_E.cumsum(dim=0) + expert_starts_E = expert_ends_E - num_tokens_per_expert_E + rows_R1 = torch.arange(x_RD.shape[0], device=x_RD.device).unsqueeze(1) + local_affinities = ( + (rows_R1 >= expert_starts_E.unsqueeze(0)) + & (rows_R1 < expert_ends_E.unsqueeze(0)) + ).to(torch.float32) + (*_, tokens, _, output) = nki_fwd( + x_RD.bfloat16(), + local_affinities, + gate_up_proj_weight, + down_proj_weight, + 1, + EP_BLOCK_SIZE, + ) + return output[:tokens].type_as(x_RD) + + +def _local(t): + """Plain local view of a (Replicate) DTensor buffer for the local MoE region.""" + return t.to_local() if isinstance(t, DTensor) else t + + +def _localize_gate_output(scores): + """Localize the router output with a Replicate local gradient. + + The experts' from_local(Replicate) input wrap already all-reduces the per-rank score + gradients; a Partial label here would sum them again (gate gradient TP x too large).""" + if isinstance(scores, DTensor): + return scores.to_local( + grad_placements=tuple(Replicate() for _ in scores.placements) + ) + return scores + + +class NKIMoE(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + num_experts: int + router: MaskedScoresTopKRouter.Config + experts: NKIGroupedExperts.Config + load_balance_coeff: float | None = 1e-3 + # As MoE.Config.load_balancing (models/common/moe.py). + load_balancing: Literal["expert_bias", "global_aux_loss"] = "expert_bias" + shared_experts: None = None + + def set_sharding_config(self, *, enable_sp: bool, enable_ep: bool) -> None: + if enable_ep: + raise NotImplementedError("Qwen3 NKI MoE supports TP-only experts") + + activation = ( + dense_sequence_parallel_placement() + if enable_sp + else dense_activation_placement(tp=spmd.I) + ) + replicated = dense_activation_placement(tp=spmd.R) + partial = dense_activation_placement(tp=spmd.P) + + self.sharding_config = ShardingConfig( + state_shardings={ + "expert_bias_E": dense_param_placement(tp=spmd.R), + "tokens_per_expert_E": dense_param_placement(tp=spmd.R), + }, + in_src_shardings={"x": activation}, + in_dst_shardings={"x": replicated}, + out_src_shardings=partial, + out_dst_shardings=activation, + ) + self.router.gate.sharding_config = ShardingConfig( + state_shardings={ + "weight": dense_param_placement(tp=spmd.R), + "bias": dense_param_placement(tp=spmd.R), + }, + in_src_shardings={"input": replicated}, + out_src_shardings=replicated, + ) + self.experts.sharding_config = ShardingConfig( + state_shardings={ + "gate_up_proj_weight": dense_param_placement(tp=spmd.S(3)), + "down_proj_weight": dense_param_placement(tp=spmd.S(1)), + }, + in_src_shardings={ + "hidden_states": replicated, + "masked_scores": replicated, + }, + in_dst_shardings={ + "hidden_states": replicated, + "masked_scores": replicated, + }, + out_src_shardings=partial, + out_dst_shardings=partial, + local_map=LocalMapConfig( + in_grad_placements=(partial, partial), + ), + ) + + def __init__(self, config: Config): + super().__init__() + self.num_experts = config.num_experts + self.router = config.router.build() + self.experts = config.experts.build() + self.load_balance_coeff = config.load_balance_coeff + self.load_balancing = config.load_balancing + if self.load_balancing not in ("expert_bias", "global_aux_loss"): + raise ValueError(f"unknown load_balancing {self.load_balancing!r}") + self.lbl_batch_mesh: DeviceMesh | None = None + self.lbl_num_microbatches: int = 1 + self.shared_experts = None + # Auxiliary-loss-free load balancing state. tokens_per_expert_E is + # accumulated in forward; expert_bias_E is updated by the + # OptimizersContainer step pre-hook (_update_expert_bias), exactly as for + # the std MoE the GPU/reference baseline runs and the DeepSeek NKI MoE. + if self.load_balance_coeff is not None: + assert self.load_balance_coeff > 0.0 + if self.load_balance_coeff is not None and self.load_balancing == "expert_bias": + self.register_buffer( + "expert_bias_E", + torch.zeros(self.num_experts, dtype=torch.float32), + persistent=True, + ) + else: + self.expert_bias_E = None + self.register_buffer( + "tokens_per_expert_E", + torch.zeros(self.num_experts, dtype=torch.float32), + persistent=False, + ) + + def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: + if buffer_device is None: + buffer_device = self.tokens_per_expert_E.device + with torch.device(buffer_device): + self.tokens_per_expert_E = torch.zeros( + self.num_experts, dtype=torch.float32 + ) + if self.expert_bias_E is not None: + self.expert_bias_E = torch.zeros(self.num_experts, dtype=torch.float32) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Run the whole MoE (router + experts) on local tensors. The wrapper's + # ShardingConfig has already redistributed x to Replicate on TP. The + # gate and experts sub-modules re-wrap this plain input through their + # own ShardingConfig (DTensor.from_local on in_src), and that wrap's + # backward all-reduces their Partial input gradients -- so the gradient + # arriving here is already fully reduced. Localize with a plain + # to_local(): declaring it Partial would sum it a second time (4x). + spec = None + if isinstance(x, DTensor): + spec = (x.device_mesh, x.placements) + x = x.to_local() + batch, sequence, dim = x.shape + flat = x.reshape(-1, dim) + lbl = ( + self.load_balancing == "global_aux_loss" + and self.load_balance_coeff is not None + and torch.is_grad_enabled() + ) + if self.router._debug_force_load_balance: + with torch.autocast(device_type=flat.device.type, dtype=torch.float32): + raw_scores = self.router.gate(flat) + raw_scores = _localize_gate_output(raw_scores) + if self.router.score_func == "sigmoid": + raw_scores = torch.sigmoid(raw_scores) + else: + raw_scores = F.softmax(raw_scores, dim=1) + masked_scores = _debug_force_load_balance_masked_scores( + raw_scores, + self.router.top_k, + self.num_experts, + self.router.route_norm, + self.router.route_scale, + ) + if lbl: + lbl_stats = ((masked_scores != 0).sum(dim=0), raw_scores.sum(dim=0)) + else: + # Pass expert_bias_E as is: when sharded it is a Replicate DTensor like + # the router's scores, and DTensor + plain Tensor is rejected. + masked_scores = self.router(flat, self.expert_bias_E, return_scores=lbl) + if lbl: + masked_scores, scores = masked_scores + masked_scores = _localize_gate_output(masked_scores) + if lbl: + # Every TP rank holds the same tokens here, so each computes the same loss term. + lbl_stats = ( + (masked_scores != 0).sum(dim=0), + # pyrefly: ignore [unbound-name] + _localize_gate_output(scores).sum(dim=0), + ) + + # Expert usage for the balancing hook (sign() makes the AC double count harmless). + with torch.no_grad(): + _local(self.tokens_per_expert_E).add_( + (masked_scores != 0).to(torch.float32).sum(dim=0) + ) + + output = self.experts(flat, masked_scores) + if isinstance(output, DTensor): + output = output.to_local() + if lbl: + output = attach_global_batch_lbl( + output, + # pyrefly: ignore [unbound-name] + *lbl_stats, + self, + top_k=self.router.top_k, + ) + output = output.reshape(batch, sequence, dim) + if spec is not None: + output = DTensor.from_local( + output, spec[0], tuple(Partial() for _ in spec[1]), run_check=False + ) + return output + + +# NKI packs Qwen3 gate/up/down expert weights under different parameter names. +# Preserve the original per-weight initialization when converting GroupedExperts. +def _default_init(tensor) -> None: + nn.init.trunc_normal_(tensor, mean=0.0, std=0.02) + + +def _gate_up_init( + w1_init: Callable | None, + w3_init: Callable | None, +) -> Callable: + def init(tensor) -> None: + (w1_init or _default_init)(tensor[:, :, 0, :]) + (w3_init or _default_init)(tensor[:, :, 1, :]) + + return init + + +def _down_init(w2_init: Callable | None) -> Callable: + def init(tensor) -> None: + (w2_init or _default_init)(tensor) + + return init + + +def to_nki_moe_config(moe: MoE.Config) -> NKIMoE.Config: + """Convert a standard Qwen3 MoE config into the Neuron NKI implementation.""" + if moe.shared_experts is not None: + raise NotImplementedError("Qwen3 NKI MoE does not support shared experts") + if not isinstance(moe.router, TokenChoiceTopKRouter.Config): + raise TypeError( + f"Expected TokenChoiceTopKRouter.Config, got {type(moe.router)!r}" + ) + if not isinstance(moe.experts, GroupedExperts.Config): + raise TypeError(f"Expected GroupedExperts.Config, got {type(moe.experts)!r}") + if ( + moe.router.num_expert_groups is not None + or moe.router.num_limited_groups is not None + ): + raise NotImplementedError( + "Qwen3 NKI MoE does not support group-limited routing" + ) + + param_init = moe.experts.param_init or {} + nki_moe = NKIMoE.Config( + num_experts=moe.num_experts, + load_balance_coeff=moe.load_balance_coeff, + load_balancing=moe.load_balancing, + router=MaskedScoresTopKRouter.Config( + gate=moe.router.gate, + top_k=moe.router.top_k, + score_func=moe.router.score_func, + route_norm=moe.router.route_norm, + route_scale=moe.router.route_scale, + _debug_force_load_balance=moe.router._debug_force_load_balance, + ), + experts=NKIGroupedExperts.Config( + dim=moe.experts.dim, + hidden_dim=moe.experts.hidden_dim, + num_experts=moe.experts.num_experts, + top_k=moe.router.top_k, + param_init={ + "gate_up_proj_weight": _gate_up_init( + param_init.get("w1_EFD"), + param_init.get("w3_EFD"), + ), + "down_proj_weight": _down_init(param_init.get("w2_EDF")), + }, + ), + ) + return nki_moe diff --git a/torchtitan_neuron/tests/README.md b/torchtitan_neuron/tests/README.md new file mode 100644 index 00000000000..b3c9027fe6f --- /dev/null +++ b/torchtitan_neuron/tests/README.md @@ -0,0 +1,23 @@ +# Testing `torchtitan_neuron` + +Tests in this directory mirror the `torchtitan_neuron` source layout. Keep +hardware-specific tests device-guarded and prefer CPU-runnable config and +registry tests when validating extension contracts. + +The test configuration disables PyTorch device-backend autoload during +collection. Tests that require Neuron packages import them explicitly after +PyTorch initialization and skip when their dependencies are unavailable. + +Run the suite from the repository root: + +```bash +pytest -s torchtitan_neuron/tests/ +``` + +`test_flex_attention.py` checks Neuron kernel lowering for static causal, +static sliding-window, and dynamic document masks. It requires `torch_neuronx` +and a Neuron device: + +```bash +pytest -s torchtitan_neuron/tests/test_flex_attention.py +``` diff --git a/torchtitan_neuron/tests/__init__.py b/torchtitan_neuron/tests/__init__.py new file mode 100644 index 00000000000..fdadd29fab5 --- /dev/null +++ b/torchtitan_neuron/tests/__init__.py @@ -0,0 +1,38 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron enablement shim for PyTorch's DTensorTestBase. + +Runs at import of the ``torchtitan_neuron.tests`` package. It must live in the +package ``__init__`` (not ``conftest.py``) because ``MultiProcessTestCase`` +spawns each rank in a fresh interpreter that re-imports the test module and its +package ``__init__`` but never ``conftest.py`` -- this is the only place the +patch reaches the ranks where the process group is built. + +torch_neuronx registers ``neuron`` as a PrivateUse1 accelerator with a ``neuron`` +c10d backend, so ``DTensorTestBase.backend`` resolves to the bare string +``"neuron"``, which ``common_dtensor.init_pg`` rejects. The compound +``"cpu:gloo,:"`` form is accepted, so return that instead. +Mirrors ``tests/unit_tests/__init__.py``; both are unnecessary once the torch +build under test accepts a PrivateUse1 device's registered backend directly. +""" + +import torch +import torch.distributed as dist +from torch.testing._internal.distributed._tensor import common_dtensor + +_PRIVATEUSE1_NAME = torch._C._get_privateuse1_backend_name() + +if not getattr(common_dtensor.DTensorTestBase, "_neuron_backend_patched", False): + _original_backend = common_dtensor.DTensorTestBase.backend + + def _backend(self): + if self.device_type == _PRIVATEUSE1_NAME: + curr = dist.get_default_backend_for_device(self.device_type) + return f"cpu:gloo,{self.device_type}:{curr}" + return _original_backend.fget(self) + + common_dtensor.DTensorTestBase.backend = property(_backend) + common_dtensor.DTensorTestBase._neuron_backend_patched = True diff --git a/torchtitan_neuron/tests/conftest.py b/torchtitan_neuron/tests/conftest.py new file mode 100644 index 00000000000..808a3134d71 --- /dev/null +++ b/torchtitan_neuron/tests/conftest.py @@ -0,0 +1,12 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +import os + + +# Import torch_neuronx explicitly from tests after torch has finished importing. +# Backend autoload can otherwise recurse through torch_neuronx while torch is +# partially initialized and duplicate-register core operator namespaces. +os.environ.setdefault("TORCH_DEVICE_BACKEND_AUTOLOAD", "0") diff --git a/torchtitan_neuron/tests/models/__init__.py b/torchtitan_neuron/tests/models/__init__.py new file mode 100644 index 00000000000..fcb62242c2d --- /dev/null +++ b/torchtitan_neuron/tests/models/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. diff --git a/torchtitan_neuron/tests/models/gpt_oss/__init__.py b/torchtitan_neuron/tests/models/gpt_oss/__init__.py new file mode 100644 index 00000000000..fcb62242c2d --- /dev/null +++ b/torchtitan_neuron/tests/models/gpt_oss/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. diff --git a/torchtitan_neuron/tests/models/gpt_oss/test_graph_trainer_config_registry.py b/torchtitan_neuron/tests/models/gpt_oss/test_graph_trainer_config_registry.py new file mode 100644 index 00000000000..10edf185eda --- /dev/null +++ b/torchtitan_neuron/tests/models/gpt_oss/test_graph_trainer_config_registry.py @@ -0,0 +1,63 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +# ruff: noqa: E402 + +import pytest + +pytest.importorskip("nkilib") +pytest.importorskip("spmd_types") +pytest.importorskip("torch_neuronx") + +from torchtitan.components.optimizer import register_moe_load_balancing_hook +from torchtitan.distributed.activation_checkpoint import FullAC +from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig + +from torchtitan_neuron.experiments.graph_trainer.gpt_oss import ( + config_registry as graph_config_registry, + model_registry, +) +from torchtitan_neuron.experiments.graph_trainer.gpt_oss.model import ( + GraphTrainerGptOssModel, +) +from torchtitan_neuron.experiments.graph_trainer.gpt_oss.parallelize import ( + parallelize_gptoss, +) +from torchtitan_neuron.models.gpt_oss.config_registry import gpt_oss_20b_nki_tp4_fsdp16 +from torchtitan_neuron.models.gpt_oss.nki_attention import NKIFlashAttentionWrapper +from torchtitan_neuron.models.gpt_oss.nki_moe import NKIMoE + + +def test_simple_fsdp_gpt_oss_neuron_config_contract(): + base = gpt_oss_20b_nki_tp4_fsdp16() + config = graph_config_registry.simple_fsdp_gpt_oss_20b_nki_tp4_fsdp16() + + assert config.model_spec.name == "simple_fsdp/gpt_oss" + assert config.model_spec.flavor == base.model_spec.flavor == "20b" + assert isinstance(config.model_spec.model, GraphTrainerGptOssModel.Config) + assert config.model_spec.parallelize_fn is parallelize_gptoss + assert config.model_spec.post_optimizer_build_fn is register_moe_load_balancing_hook + + assert isinstance(config.compile, GraphTrainerCompileConfig) + assert config.compile.enable is True + assert config.compile.backend == "neuron" + assert config.compile.mode == "jit" + assert config.compile.components == ["model", "loss", "grad_clip"] + assert isinstance(config.activation_checkpoint, FullAC.Config) + + for layer in config.model_spec.model.layers: + assert isinstance(layer.moe, NKIMoE.Config) + assert isinstance( + layer.attention.inner_attention, + NKIFlashAttentionWrapper.Config, + ) + + +def test_gpt_oss_model_registry_is_neuron_owned(): + model_spec = model_registry("20b") + + assert model_spec.name == "graph_trainer/gpt_oss" + assert isinstance(model_spec.model, GraphTrainerGptOssModel.Config) + assert model_spec.parallelize_fn is parallelize_gptoss diff --git a/torchtitan_neuron/tests/models/gpt_oss/test_trainer_config_registry.py b/torchtitan_neuron/tests/models/gpt_oss/test_trainer_config_registry.py new file mode 100644 index 00000000000..c3e31e5372d --- /dev/null +++ b/torchtitan_neuron/tests/models/gpt_oss/test_trainer_config_registry.py @@ -0,0 +1,83 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +# ruff: noqa: E402 + +import pytest +import torch + +pytest.importorskip("nkilib") +pytest.importorskip("spmd_types") +pytest.importorskip("torch_neuronx") + +from torchtitan.models.gpt_oss import model_registry + +from torchtitan_neuron.models.gpt_oss.config_registry import gpt_oss_20b_nki_tp4_fsdp16 +from torchtitan_neuron.models.gpt_oss.nki_attention import NKIFlashAttentionWrapper +from torchtitan_neuron.models.gpt_oss.nki_moe import NKIMoE, to_nki_moe_config + +_PARAM_SHAPES = { + "mlp1_weight": (4, 8, 2, 6), + "mlp1_bias": (4, 2, 6), + "mlp2_weight": (4, 6, 8), + "mlp2_bias": (4, 8), +} +_STOCK_NAMES = { + "mlp1_weight": "mlp1_weight_EGD", + "mlp1_bias": "mlp1_bias_EG", + "mlp2_weight": "mlp2_weight_EDF", + "mlp2_bias": "mlp2_bias_ED", +} + + +def _init(fn, shape, seed): + torch.manual_seed(seed) + tensor = torch.empty(shape) + fn(tensor) + return tensor + + +def test_to_nki_moe_config_keeps_stock_router_and_init(): + stock = model_registry("20b").model.layers[3].moe + nki = to_nki_moe_config(stock) + + assert isinstance(nki, NKIMoE.Config) + assert nki.num_experts == stock.num_experts + assert nki.load_balance_coeff == stock.load_balance_coeff + assert nki.router.top_k == nki.experts.top_k == stock.router.top_k + assert nki.router.gate == stock.router.gate + assert nki.experts.swiglu_limit == stock.experts.swiglu_limit + + for name, shape in _PARAM_SHAPES.items(): + stock_init = stock.experts.param_init[_STOCK_NAMES[name]] + expected = _init(stock_init, shape, seed=7) + if name == "mlp1_bias": + # GPT-OSS adds +1 to the up bias (index 1 of the packed gate/up axis). + expected[:, 1, :] += 1.0 + torch.testing.assert_close( + _init(nki.experts.param_init[name], shape, seed=7), expected + ) + + +def test_gpt_oss_nki_recipe_uses_nki_modules(): + config = gpt_oss_20b_nki_tp4_fsdp16() + model = config.model_spec.model + model.update_from_config(config=config) + + for layer in model.layers: + assert isinstance(layer.moe, NKIMoE.Config) + assert isinstance( + layer.attention.inner_attention, NKIFlashAttentionWrapper.Config + ) + assert layer.attention.rope.use_prefix_cache_for_positions is True + # Placements come from NKIFlashAttentionWrapper.Config.set_sharding_config. + assert set( + layer.attention.inner_attention.sharding_config.in_src_shardings + ) == { + "q", + "k", + "v", + "sinks", + } diff --git a/torchtitan_neuron/tests/models/llama3/__init__.py b/torchtitan_neuron/tests/models/llama3/__init__.py new file mode 100644 index 00000000000..fcb62242c2d --- /dev/null +++ b/torchtitan_neuron/tests/models/llama3/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. diff --git a/torchtitan_neuron/tests/models/llama3/test_graph_trainer_config_registry.py b/torchtitan_neuron/tests/models/llama3/test_graph_trainer_config_registry.py new file mode 100644 index 00000000000..9535fe74a0b --- /dev/null +++ b/torchtitan_neuron/tests/models/llama3/test_graph_trainer_config_registry.py @@ -0,0 +1,69 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +import pytest + + +pytest.importorskip("spmd_types") + +from torchtitan.experiments.graph_trainer.configs import ( # noqa: E402 + GraphTrainerCompileConfig, +) +from torchtitan.experiments.graph_trainer.llama3.model import ( # noqa: E402 + GraphTrainerLlama3Model, +) +from torchtitan.models.common.attention import ScaledDotProductAttention # noqa: E402 +from torchtitan.models.llama3.config_registry import llama3_8b # noqa: E402 + +from torchtitan_neuron.experiments.graph_trainer.llama3 import ( # noqa: E402 + config_registry as graph_trainer_llama3_config_registry, +) + + +def test_simple_fsdp_llama3_8b_neuron_override_config_contract(): + base = llama3_8b() + config = graph_trainer_llama3_config_registry.simple_fsdp_llama3_8b() + + assert config.model_spec.name == "simple_fsdp/llama3" + assert config.model_spec.flavor == base.model_spec.flavor == "8B" + assert isinstance(config.model_spec.model, GraphTrainerLlama3Model.Config) + assert len(config.model_spec.model.layers) == len(base.model_spec.model.layers) + assert isinstance( + config.model_spec.model.layers[0].attention.inner_attention, + ScaledDotProductAttention.Config, + ) + + assert isinstance(config.compile, GraphTrainerCompileConfig) + assert config.compile.enable is True + assert config.compile.backend == "neuron" + assert config.compile.mode == "jit" + assert config.compile.components == ["model", "loss", "grad_clip"] + + assert config.override.imports == [] + assert all( + layer.attention.rope.use_prefix_cache_for_positions + for layer in config.model_spec.model.layers + ) + + +@pytest.mark.parametrize( + "config_factory", + [ + graph_trainer_llama3_config_registry.simple_fsdp_llama3_8b, + graph_trainer_llama3_config_registry.simple_fsdp_llama3_8b_fullac, + graph_trainer_llama3_config_registry.simple_fsdp_llama3_70b_fullac, + ], +) +def test_simple_fsdp_llama3_uses_sdpa_backend(config_factory): + config = config_factory() + + assert config.model_spec is not None + assert all( + isinstance( + layer.attention.inner_attention, + ScaledDotProductAttention.Config, + ) + for layer in config.model_spec.model.layers + ) diff --git a/torchtitan_neuron/tests/models/llama3/test_trainer_config_registry.py b/torchtitan_neuron/tests/models/llama3/test_trainer_config_registry.py new file mode 100644 index 00000000000..2a971af8d89 --- /dev/null +++ b/torchtitan_neuron/tests/models/llama3/test_trainer_config_registry.py @@ -0,0 +1,54 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +import pytest + +pytest.importorskip("spmd_types") + +from torchtitan.distributed.activation_checkpoint import ( # noqa: E402 + FullAC, + SelectiveAC, +) +from torchtitan.models.common.attention import ScaledDotProductAttention # noqa: E402 + +from torchtitan_neuron.models.llama3 import config_registry # noqa: E402 + +# (recipe, full_ac, compile, n_layers) +_RECIPES = [ + (config_registry.llama3_8b_sac_eager, False, False, 32), + (config_registry.llama3_8b_fullac_eager, True, False, 32), + (config_registry.llama3_70b_fullac_eager, True, False, 80), + (config_registry.llama3_8b_sac_compile, False, True, 32), + (config_registry.llama3_8b_fullac_compile, True, True, 32), + (config_registry.llama3_70b_fullac_compile, True, True, 80), + (config_registry.llama3_70b_16layers_fullac_compile, True, True, 16), +] + + +@pytest.mark.parametrize( + ("recipe", "full_ac", "compile", "n_layers"), + _RECIPES, + ids=[r[0].__name__ for r in _RECIPES], +) +def test_neuron_dense_recipe_contract(recipe, full_ac, compile, n_layers): + config = recipe() + layers = config.model_spec.model.layers + + assert len(layers) == n_layers + for layer in layers: + assert isinstance( + layer.attention.inner_attention, ScaledDotProductAttention.Config + ) + assert layer.attention.rope.use_prefix_cache_for_positions is True + assert isinstance( + config.activation_checkpoint, FullAC.Config if full_ac else SelectiveAC.Config + ) + assert config.dataloader.dataset == "c4_test" + assert config.metrics.log_freq == 1 + assert config.compile.enable is compile + if compile: + assert config.compile.backend == "neuron" + assert config.compile.components == ["model", "loss", "grad_clip"] + assert config.override.imports == [] diff --git a/torchtitan_neuron/tests/models/qwen3/__init__.py b/torchtitan_neuron/tests/models/qwen3/__init__.py new file mode 100644 index 00000000000..fcb62242c2d --- /dev/null +++ b/torchtitan_neuron/tests/models/qwen3/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. diff --git a/torchtitan_neuron/tests/models/qwen3/test_global_lbl.py b/torchtitan_neuron/tests/models/qwen3/test_global_lbl.py new file mode 100644 index 00000000000..7d690057ac7 --- /dev/null +++ b/torchtitan_neuron/tests/models/qwen3/test_global_lbl.py @@ -0,0 +1,103 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +# ruff: noqa: E402 +"""Qwen3 global-batch load-balancing loss on the NKI MoE.""" + +import copy +from unittest import mock + +import pytest +import torch + +pytest.importorskip("nkilib") +pytest.importorskip("spmd_types") +pytest.importorskip("torch_neuronx") + +from torch.distributed.tensor import DTensor +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) +from torchtitan.distributed.parallel_dims import ParallelDims +from torchtitan.models.common.nn_modules import Linear + +from torchtitan_neuron.models.qwen3 import nki_moe +from torchtitan_neuron.tests.models.qwen3.test_nki_moe import ( + _ROUTER_GRAD_BATCH as _BATCH, + _ROUTER_GRAD_DIM as _DIM, + _router_grad_fake_nki_fwd as _fake_nki_fwd, + _ROUTER_GRAD_HIDDEN_DIM as _HIDDEN_DIM, + _ROUTER_GRAD_NUM_EXPERTS as _NUM_EXPERTS, + _ROUTER_GRAD_SEQ as _SEQ, + _ROUTER_GRAD_TOP_K as _TOP_K, +) + + +def _nki_moe(): + config = nki_moe.NKIMoE.Config( + num_experts=_NUM_EXPERTS, + load_balance_coeff=1.0, + load_balancing="global_aux_loss", + router=nki_moe.MaskedScoresTopKRouter.Config( + gate=Linear.Config(in_features=_DIM, out_features=_NUM_EXPERTS, bias=False), + top_k=_TOP_K, + score_func="softmax", + route_norm=True, + ), + experts=nki_moe.NKIGroupedExperts.Config( + dim=_DIM, hidden_dim=_HIDDEN_DIM, num_experts=_NUM_EXPERTS, top_k=_TOP_K + ), + ) + config.set_sharding_config(enable_sp=False, enable_ep=False) + torch.manual_seed(0) + moe = config.build() + with torch.no_grad(): + for param in moe.parameters(): + param.normal_(std=0.5) + return moe.train() + + +def _nki_gate_grad(moe, x_BSD, main_weight): + with mock.patch.object(nki_moe, "nki_fwd", _fake_nki_fwd([])): + out = moe(x_BSD) + out = out.full_tensor() if isinstance(out, DTensor) else out + (main_weight * out.float().pow(2).sum()).backward() + grad = moe.router.gate.weight.grad + return grad.full_tensor() if isinstance(grad, DTensor) else grad + + +class TestNKIUnderTP(DTensorTestBase): + @property + def world_size(self) -> int: + return 2 + + @property + def device_type(self) -> str: + return "cpu" + + @with_comms + def test_gate_gradient_matches_tp1(self): + parallel_dims = ParallelDims( + dp_replicate=1, + dp_shard=1, + cp=1, + tp=self.world_size, + pp=1, + ep=1, + world_size=self.world_size, + ) + with mock.patch("torchtitan.distributed.parallel_dims.device_type", "cpu"): + parallel_dims.build_mesh() + torch.manual_seed(1) + x_BSD = torch.randn(_BATCH, _SEQ, _DIM) + # Aux term alone (zero main loss), then with a main loss: both must equal TP1. + for main_weight in (0.0, 1.0 / (_BATCH * _SEQ)): + moe = _nki_moe() + reference = copy.deepcopy(moe) + moe.parallelize(parallel_dims) + want = _nki_gate_grad(reference, x_BSD, main_weight) + assert float(want.norm()) > 0 + torch.testing.assert_close(_nki_gate_grad(moe, x_BSD, main_weight), want) diff --git a/torchtitan_neuron/tests/models/qwen3/test_graph_trainer_config_registry.py b/torchtitan_neuron/tests/models/qwen3/test_graph_trainer_config_registry.py new file mode 100644 index 00000000000..e3d3e1a6348 --- /dev/null +++ b/torchtitan_neuron/tests/models/qwen3/test_graph_trainer_config_registry.py @@ -0,0 +1,198 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +# ruff: noqa: E402 + +import inspect + +import pytest + +pytest.importorskip("nkilib") +pytest.importorskip("spmd_types") +pytest.importorskip("torch_neuronx") + +from torchtitan.components.optimizer import register_moe_load_balancing_hook +from torchtitan.config import CompileConfig +from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC +from torchtitan.distributed.parallel_dims import ParallelDims +from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig +from torchtitan.experiments.graph_trainer.qwen3 import ( + config_registry as core_graph_config_registry, +) +from torchtitan.experiments.graph_trainer.qwen3.parallelize import parallelize_qwen3 + +from torchtitan_neuron.experiments.graph_trainer.qwen3 import ( + config_registry as graph_config_registry, +) +from torchtitan_neuron.models.qwen3 import ( + config_registry as trainer_config_registry, + nki_moe, +) + + +def _config_factories(): + factories = [] + for name, factory in inspect.getmembers( + graph_config_registry, + inspect.isfunction, + ): + if ( + not name.startswith("qwen3_") + or factory.__module__ != graph_config_registry.__name__ + ): + continue + if any( + parameter.default is inspect.Parameter.empty + for parameter in inspect.signature(factory).parameters.values() + ): + continue + factories.append(pytest.param(factory, id=name)) + return factories + + +@pytest.mark.parametrize("config_factory", _config_factories()) +def test_graph_trainer_recipes_finalize_model_sharding(config_factory): + config = config_factory() + + config.model_spec.model.update_from_config(config=config) + + for layer in config.model_spec.model.layers: + moe = layer.moe + if isinstance(moe, nki_moe.NKIMoE.Config): + assert moe.sharding_config is not None + assert moe.router.gate.sharding_config is not None + assert moe.experts.sharding_config is not None + + +def test_neuron_graph_trainer_recipes_are_not_registered_in_core(): + assert not hasattr( + core_graph_config_registry, + "graph_trainer_qwen3_235b_scaled", + ) + assert not hasattr( + core_graph_config_registry, + "simple_fsdp_qwen3_32b_tp4_fsdp16_fullac", + ) + + +def test_relocated_235b_aot_graph_trainer_recipes(): + for config in ( + graph_config_registry.graph_trainer_qwen3_235b_scaled(), + graph_config_registry.graph_trainer_qwen3_235b_4layers_ep64_nle1_tp4_neuron_ep(), + ): + assert config.model_spec.name == "qwen3" + assert isinstance(config.compile, GraphTrainerCompileConfig) + assert config.compile.enable is True + assert config.compile.mode == "aot_fx_trace" + + +def test_qwen3_30b_simple_fsdp_recipe_reuses_trainer_model(): + base = trainer_config_registry.qwen3_nki_moe_config( + "30B-A3B", + compile_config=CompileConfig(enable=False, backend="neuron"), + ) + config = graph_config_registry.qwen3_30b_a3b_nki_fullac_simplefsdp() + + assert config.model_spec.flavor == base.model_spec.flavor + assert len(config.model_spec.model.layers) == 48 + assert config.training.local_batch_size == base.training.local_batch_size + assert config.training.steps == base.training.steps + assert config.dataloader.dataset == base.dataloader.dataset + assert config.model_spec.name == "simple_fsdp/qwen3" + assert isinstance(config.compile, GraphTrainerCompileConfig) + assert config.compile.enable is True + assert config.compile.mode == "jit" + assert config.compile.components == ["model", "loss", "grad_clip"] + assert config.model_spec.parallelize_fn is parallelize_qwen3 + assert config.model_spec.post_optimizer_build_fn is register_moe_load_balancing_hook + assert config.override.imports == [] + + +def test_qwen3_235b_six_layer_4k_neuron_ep_recipes(): + trainer = ( + trainer_config_registry.qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_compile() + ) + graph = ( + graph_config_registry.qwen3_235b_6layers_ep64_nle1_tp4_neuron_ep_compile_simplefsdp() + ) + + for config in (trainer, graph): + assert len(config.model_spec.model.layers) == 6 + assert config.training.seq_len == 4096 + assert isinstance(config.activation_checkpoint, SelectiveAC.Config) + assert config.parallelism.expert_parallel_degree == 64 + assert config.parallelism.tensor_parallel_degree == 4 + assert config.compile.enable is True + assert config.compile.backend == "neuron" + assert config.compile.components == ["model", "loss"] + for layer in config.model_spec.model.layers: + assert layer.moe.num_experts == 64 + assert layer.moe.router.top_k == 8 + assert isinstance(layer.moe.experts, nki_moe.BWMMGroupedExperts.Config) + + assert graph.model_spec.name == "simple_fsdp/qwen3" + + +def test_qwen3_235b_ep128_4k_neuron_ep_recipes(): + trainer = ( + trainer_config_registry.qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_compile() + ) + graph = ( + graph_config_registry.qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_compile_simplefsdp() + ) + + for config in (trainer, graph): + assert len(config.model_spec.model.layers) == 6 + assert config.training.seq_len == 4096 + assert isinstance(config.activation_checkpoint, SelectiveAC.Config) + assert config.parallelism.expert_parallel_degree == 128 + assert config.parallelism.tensor_parallel_degree == 4 + for layer in config.model_spec.model.layers: + assert layer.moe.num_experts == 128 + assert layer.moe.router.top_k == 8 + assert isinstance(layer.moe.experts, nki_moe.BWMMGroupedExperts.Config) + + +def test_qwen3_235b_94layer_ep128_fsdp64_pp1_topology(): + config = ( + trainer_config_registry.qwen3_235b_94layers_ep128_tp4_fsdp64_pp1_4k_neuron_ep_compile() + ) + + assert len(config.model_spec.model.layers) == 94 + assert config.training.seq_len == 4096 + assert config.training.local_batch_size == 1 + assert isinstance(config.activation_checkpoint, FullAC.Config) + + parallel_dims = ParallelDims.from_config(config.parallelism, world_size=256) + assert parallel_dims.pp == 1 + assert parallel_dims.dp_replicate == 1 + assert parallel_dims.dp_shard == 64 + assert parallel_dims.tp == 4 + assert parallel_dims.ep == 128 + expert_fsdp = ( + parallel_dims.dp_shard * parallel_dims.cp * parallel_dims.tp // parallel_dims.ep + ) + assert expert_fsdp == 2 + + +def test_qwen3_235b_ep128_pp2_preserves_pipeline_configuration(): + trainer = ( + trainer_config_registry.qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_pp2_compile() + ) + graph = ( + graph_config_registry.qwen3_235b_6layers_ep128_nle1_tp4_neuron_ep_pp2_compile_simplefsdp() + ) + + for config in (trainer, graph): + assert config.parallelism.expert_parallel_degree == 128 + assert config.parallelism.tensor_parallel_degree == 4 + assert config.parallelism.pipeline_parallel_degree == 2 + assert config.parallelism.pipeline_parallel_schedule == "1F1B" + assert config.parallelism.pipeline_parallel_first_stage_less_layers == 0 + assert config.parallelism.pipeline_parallel_last_stage_less_layers == 0 + assert config.training.local_batch_size == 16 + + assert graph.model_spec.name == "simple_fsdp/qwen3" + assert graph.model_spec.pipelining_fn is not None diff --git a/torchtitan_neuron/tests/models/qwen3/test_nki_moe.py b/torchtitan_neuron/tests/models/qwen3/test_nki_moe.py new file mode 100644 index 00000000000..44860cbe41a --- /dev/null +++ b/torchtitan_neuron/tests/models/qwen3/test_nki_moe.py @@ -0,0 +1,483 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +# ruff: noqa: E402 + +import contextlib +import copy +from unittest import mock + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + +pytest.importorskip("nkilib") +pytest.importorskip("spmd_types") +pytest.importorskip("torch_neuronx") + +from torch.distributed.tensor import DTensor, Partial, Replicate +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) +from torchtitan.distributed.parallel_dims import ParallelDims +from torchtitan.models.common.moe import GroupedExperts, MoE, TokenChoiceTopKRouter +from torchtitan.models.common.nn_modules import Linear +from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher + +from torchtitan_neuron.models.qwen3 import nki_moe + + +def _fake_nki_fwd( + hidden_states, + masked_scores, + gate_up_proj_weight, + down_proj_weight, + top_k, + block_size, +): + del gate_up_proj_weight, down_proj_weight, top_k, block_size + output = hidden_states * masked_scores.sum(dim=-1, keepdim=True).to( + hidden_states.dtype + ) + return ( + masked_scores.reshape(-1, 1), + torch.empty(0, device=hidden_states.device, dtype=hidden_states.dtype), + torch.empty(0, device=hidden_states.device, dtype=hidden_states.dtype), + torch.empty(0, device=hidden_states.device, dtype=torch.int64), + torch.empty(0, device=hidden_states.device, dtype=torch.int64), + hidden_states.shape[0], + masked_scores.shape[1], + output, + ) + + +def _tiny_nki_moe_config(): + dim = 8 + num_experts = 4 + top_k = 2 + return nki_moe.NKIMoE.Config( + num_experts=num_experts, + load_balance_coeff=None, + router=nki_moe.MaskedScoresTopKRouter.Config( + gate=Linear.Config( + in_features=dim, + out_features=num_experts, + bias=False, + ), + top_k=top_k, + score_func="sigmoid", + route_norm=True, + ), + experts=nki_moe.NKIGroupedExperts.Config( + dim=dim, + hidden_dim=16, + num_experts=num_experts, + top_k=top_k, + ), + ) + + +def _stock_moe_config(): + num_experts = 4 + top_k = 2 + return MoE.Config( + num_experts=num_experts, + load_balance_coeff=1e-3, + shared_experts=None, + router=TokenChoiceTopKRouter.Config( + num_experts=num_experts, + gate=Linear.Config(in_features=8, out_features=num_experts, bias=False), + top_k=top_k, + score_func="sigmoid", + route_norm=True, + ), + experts=GroupedExperts.Config( + dim=8, + hidden_dim=16, + num_experts=num_experts, + token_dispatcher=LocalTokenDispatcher.Config( + num_experts=num_experts, + top_k=top_k, + ), + ), + ) + + +def test_masked_scores_router_executes_and_keeps_top_k_scores(): + torch.manual_seed(0) + router = nki_moe.MaskedScoresTopKRouter( + nki_moe.MaskedScoresTopKRouter.Config( + gate=Linear.Config(in_features=8, out_features=4, bias=False), + top_k=2, + score_func="sigmoid", + route_norm=True, + ) + ) + + scores = router(torch.randn(3, 8)) + + assert scores.shape == (3, 4) + assert torch.count_nonzero(scores, dim=-1).tolist() == [2, 2, 2] + torch.testing.assert_close(scores.sum(dim=-1), torch.ones(3)) + + +def test_nki_moe_initializes_compiles_and_executes_with_mocked_nki_op(monkeypatch): + torch.manual_seed(0) + monkeypatch.setattr(nki_moe, "nki_fwd", _fake_nki_fwd) + model = nki_moe.NKIMoE(_tiny_nki_moe_config()) + + x = torch.randn(2, 3, 8) + eager_output = model(x) + compiled_output = torch.compile(model, backend="eager", fullgraph=True)(x) + + assert eager_output.shape == x.shape + torch.testing.assert_close(compiled_output, eager_output) + + +def test_bwmm_grouped_experts_packs_single_expert_weights(monkeypatch): + captured = {} + + def fake_nki_fwd( + hidden_states, + masked_scores, + gate_up_proj_weight, + down_proj_weight, + top_k, + block_size, + ): + captured["masked_scores"] = masked_scores + captured["gate_up_shape"] = gate_up_proj_weight.shape + captured["down_shape"] = down_proj_weight.shape + captured["top_k"] = top_k + captured["block_size"] = block_size + gate = hidden_states @ gate_up_proj_weight[0, :, 0, :] + up = hidden_states @ gate_up_proj_weight[0, :, 1, :] + output = (torch.nn.functional.silu(gate) * up) @ down_proj_weight[0] + return ( + masked_scores.reshape(-1, 1), + torch.empty(0), + torch.empty(0), + torch.empty(0), + torch.empty(0), + hidden_states.shape[0], + 1, + output, + ) + + monkeypatch.setattr(nki_moe, "nki_fwd", fake_nki_fwd) + torch.manual_seed(0) + experts = nki_moe.BWMMGroupedExperts( + nki_moe.BWMMGroupedExperts.Config( + dim=8, + hidden_dim=16, + num_experts=1, + token_dispatcher=LocalTokenDispatcher.Config( + num_experts=1, + top_k=1, + ), + ) + ) + with torch.no_grad(): + experts.w1_EFD.normal_(std=0.02) + experts.w2_EDF.normal_(std=0.02) + experts.w3_EFD.normal_(std=0.02) + x = torch.randn(32, 8, dtype=torch.bfloat16, requires_grad=True) + output = experts._experts_forward(x, torch.tensor([32], dtype=torch.int32)) + expected = ( + torch.nn.functional.silu(x @ experts.w1_EFD[0].bfloat16().T) + * (x @ experts.w3_EFD[0].bfloat16().T) + ) @ experts.w2_EDF[0].bfloat16().T + + assert captured["gate_up_shape"] == (1, 8, 2, 16) + assert captured["down_shape"] == (1, 16, 8) + assert captured["top_k"] == 1 + assert captured["block_size"] == nki_moe.EP_BLOCK_SIZE + torch.testing.assert_close(captured["masked_scores"], torch.ones(32, 1)) + torch.testing.assert_close(output, expected) + + output.float().sum().backward() + assert x.grad is not None + assert experts.w1_EFD.grad is not None + assert experts.w2_EDF.grad is not None + assert experts.w3_EFD.grad is not None + + +def test_single_expert_block_metadata_skips_only_padding(): + tokens = nki_moe.BLOCK_SIZE + 7 + num_blocks = 3 + token_ids, expert_ids = nki_moe._single_expert_block_metadata( + tokens, + num_blocks, + nki_moe.EP_BLOCK_SIZE, + torch.device("cpu"), + ) + + torch.testing.assert_close( + token_ids[:tokens], + torch.arange(tokens, dtype=torch.int32), + ) + torch.testing.assert_close( + token_ids[tokens:], + torch.full( + (num_blocks * nki_moe.EP_BLOCK_SIZE - tokens,), + -1, + dtype=torch.int32, + ), + ) + torch.testing.assert_close( + expert_ids, + torch.zeros((num_blocks, 1), dtype=torch.int32), + ) + + +def test_qwen3_nki_conversion_does_not_preempt_decoder_sharding(): + stock_moe = _stock_moe_config() + + replacement = nki_moe.to_nki_moe_config(stock_moe) + + assert isinstance(replacement, nki_moe.NKIMoE.Config) + assert replacement.num_experts == stock_moe.num_experts + assert replacement.load_balance_coeff == stock_moe.load_balance_coeff + assert isinstance(replacement.router, nki_moe.MaskedScoresTopKRouter.Config) + assert isinstance(replacement.experts, nki_moe.NKIGroupedExperts.Config) + assert replacement.router.top_k == stock_moe.router.top_k + assert replacement.experts.top_k == stock_moe.router.top_k + assert replacement.sharding_config is None + assert replacement.router.gate.sharding_config is None + assert replacement.experts.sharding_config is None + + +_ROUTER_GRAD_DIM = 8 +_ROUTER_GRAD_HIDDEN_DIM = 16 +_ROUTER_GRAD_NUM_EXPERTS = 4 +_ROUTER_GRAD_TOP_K = 2 +_ROUTER_GRAD_BATCH = 2 +_ROUTER_GRAD_SEQ = 8 + + +def _router_grad_fake_nki_fwd(score_grads: list[torch.Tensor]): + def fake_nki_fwd( + hidden_states, + masked_scores, + gate_up_proj_weight, + down_proj_weight, + top_k, + block_size, + ): + del top_k, block_size + masked_scores.register_hook(score_grads.append) + gate = torch.einsum( + "bd,edh->beh", + hidden_states, + gate_up_proj_weight[:, :, 0, :], + ) + up = torch.einsum( + "bd,edh->beh", + hidden_states, + gate_up_proj_weight[:, :, 1, :], + ) + expert_output = torch.einsum( + "beh,ehd->bed", + F.silu(gate) * up, + down_proj_weight, + ) + output = torch.einsum("be,bed->bd", masked_scores, expert_output) + tokens, num_experts = masked_scores.shape + return (None,) * 5 + (tokens, num_experts, output) + + return fake_nki_fwd + + +def _partial_gate_output(scores): + if isinstance(scores, DTensor): + return scores.to_local( + grad_placements=tuple(Partial() for _ in scores.placements) + ) + return scores + + +def _router_grad_nki_moe(*, enable_sp: bool): + config = nki_moe.NKIMoE.Config( + num_experts=_ROUTER_GRAD_NUM_EXPERTS, + load_balance_coeff=None, + router=nki_moe.MaskedScoresTopKRouter.Config( + gate=Linear.Config( + in_features=_ROUTER_GRAD_DIM, + out_features=_ROUTER_GRAD_NUM_EXPERTS, + bias=False, + ), + top_k=_ROUTER_GRAD_TOP_K, + score_func="softmax", + route_norm=True, + ), + experts=nki_moe.NKIGroupedExperts.Config( + dim=_ROUTER_GRAD_DIM, + hidden_dim=_ROUTER_GRAD_HIDDEN_DIM, + num_experts=_ROUTER_GRAD_NUM_EXPERTS, + top_k=_ROUTER_GRAD_TOP_K, + ), + ) + config.set_sharding_config(enable_sp=enable_sp, enable_ep=False) + torch.manual_seed(0) + moe = config.build() + with torch.no_grad(): + for parameter in moe.parameters(): + parameter.normal_(std=0.5) + return moe + + +def _router_grad_forward_backward(moe, inputs, output_grad, score_grads): + with mock.patch.object( + nki_moe, + "nki_fwd", + _router_grad_fake_nki_fwd(score_grads), + ): + output = moe(inputs) + if isinstance(output, DTensor): + output = output.full_tensor() + (output * output_grad).sum().backward() + return moe.router.gate.weight.grad + + +def _all_gather(tensor): + gathered = [torch.empty_like(tensor) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, tensor.contiguous()) + return gathered + + +class TestNKIMoERouterGradUnderTP(DTensorTestBase): + @property + def world_size(self) -> int: + return 2 + + @property + def device_type(self) -> str: + return "cpu" + + def _gate_grads(self, *, enable_sp: bool, localize=None): + parallel_dims = ParallelDims( + dp_replicate=1, + dp_shard=1, + cp=1, + tp=self.world_size, + pp=1, + ep=1, + world_size=self.world_size, + ) + with mock.patch( + "torchtitan.distributed.parallel_dims.device_type", + self.device_type, + ): + parallel_dims.build_mesh() + + moe = _router_grad_nki_moe(enable_sp=enable_sp) + reference = copy.deepcopy(moe) + moe.parallelize(parallel_dims) + + torch.manual_seed(1) + inputs = torch.randn( + _ROUTER_GRAD_BATCH, + _ROUTER_GRAD_SEQ, + _ROUTER_GRAD_DIM, + ) + output_grad = torch.randn( + _ROUTER_GRAD_BATCH, + _ROUTER_GRAD_SEQ, + _ROUTER_GRAD_DIM, + ) + local_inputs = ( + inputs.chunk(self.world_size, dim=1)[self.rank].contiguous() + if enable_sp + else inputs + ) + + reference_score_grads = [] + tp_score_grads = [] + reference_grad = _router_grad_forward_backward( + reference, + inputs, + output_grad, + reference_score_grads, + ) + label = ( + mock.patch.object(nki_moe, "_localize_gate_output", localize) + if localize is not None + else contextlib.nullcontext() + ) + with label: + tp_grad = _router_grad_forward_backward( + moe, + local_inputs, + output_grad, + tp_score_grads, + ) + + assert len(reference_score_grads) == len(tp_score_grads) == 1 + assert float(reference_grad.norm()) > 0.1 + return ( + tp_grad, + tp_score_grads[0], + reference_grad, + reference_score_grads[0], + ) + + @with_comms + def test_replicate_label_matches_the_tp1_reference(self): + for enable_sp in (False, True): + ( + tp_grad, + tp_score_grad, + reference_grad, + reference_score_grad, + ) = self._gate_grads(enable_sp=enable_sp) + + shares = _all_gather(tp_score_grad) + assert not all(torch.equal(share, shares[0]) for share in shares[1:]) + torch.testing.assert_close( + torch.stack(shares).sum(dim=0), + reference_score_grad, + ) + + assert isinstance(tp_grad, DTensor) + assert tuple(tp_grad.placements) == (Replicate(),) + local_grad = tp_grad.to_local() + ratio = float(local_grad.norm() / reference_grad.norm()) + torch.testing.assert_close( + local_grad, + reference_grad, + msg=( + f"TP{self.world_size} sp={enable_sp}: router gate gradient " + f"is {ratio:.6f}x its TP=1 value" + ), + ) + for other in _all_gather(local_grad): + assert torch.equal(other, local_grad) + + @with_comms + def test_partial_label_scales_the_gradient_by_the_tp_degree(self): + for enable_sp in (False, True): + tp_grad, _, reference_grad, _ = self._gate_grads( + enable_sp=enable_sp, + localize=_partial_gate_output, + ) + effective_grad = tp_grad.full_tensor() + ratio = float(effective_grad.norm() / reference_grad.norm()) + torch.testing.assert_close( + effective_grad, + self.world_size * reference_grad, + msg=( + f"TP{self.world_size} sp={enable_sp}: Partial scaled the " + f"router gate gradient by {ratio:.6f}x, not " + f"{self.world_size}x" + ), + ) + + +class TestNKIMoERouterGradUnderTP4(TestNKIMoERouterGradUnderTP): + @property + def world_size(self) -> int: + return 4 diff --git a/torchtitan_neuron/tests/models/qwen3/test_trainer_config_registry.py b/torchtitan_neuron/tests/models/qwen3/test_trainer_config_registry.py new file mode 100644 index 00000000000..f87061f6c6e --- /dev/null +++ b/torchtitan_neuron/tests/models/qwen3/test_trainer_config_registry.py @@ -0,0 +1,237 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +import inspect + +import pytest + +pytest.importorskip("nkilib") +pytest.importorskip("spmd_types") +pytest.importorskip("torch_neuronx") + +from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss +from torchtitan.distributed.activation_checkpoint import FullAC +from torchtitan.models.common.attention import ScaledDotProductAttention +from torchtitan.models.common.rope import CosSinRoPE +from torchtitan.models.qwen3 import ( + config_registry as core_config_registry, + model_registry as core_model_registry, + qwen3_configs, +) + +from torchtitan_neuron.components.quantization import torchao_supports_neuron +from torchtitan_neuron.models.qwen3 import config_registry, model_registry, nki_moe + + +def _config_factories(): + factories = [] + for name, factory in inspect.getmembers(config_registry, inspect.isfunction): + if ( + not name.startswith("qwen3_") + or factory.__module__ != config_registry.__name__ + ): + continue + if any( + parameter.default is inspect.Parameter.empty + for parameter in inspect.signature(factory).parameters.values() + ): + continue + # MXFP8 recipes attach a torchao-backed converter at build time; skip them + # unless torchao with Neuron support is installed (public torchao lacks it). + marks = ( + pytest.mark.skipif( + not torchao_supports_neuron(), + reason="MXFP8 configs require torchao with Neuron support", + ) + if "mxfp8" in name + else () + ) + factories.append(pytest.param(factory, id=name, marks=marks)) + return factories + + +@pytest.fixture +def recipe_config(config_factory): + if "mxfp8" in config_factory.__name__: + pytest.importorskip("torchao") + return config_factory() + + +@pytest.mark.parametrize("config_factory", _config_factories()) +def test_trainer_recipes_finalize_model_sharding(recipe_config): + config = recipe_config + + config.model_spec.model.update_from_config(config=config) + + for layer in config.model_spec.model.layers: + moe = layer.moe + if isinstance(moe, nki_moe.NKIMoE.Config): + assert moe.sharding_config is not None + assert moe.router.gate.sharding_config is not None + assert moe.experts.sharding_config is not None + + +def test_model_map_contains_moe_presets_only(): + assert set(config_registry._QWEN3_NKI_MOE_PRESETS) == { + "30B-A3B", + "235B-A22B", + } + + +def test_qwen3_235b_base_recipe_is_neuron_owned_and_importable(): + assert not hasattr(core_config_registry, "qwen3_235b_a22b") + + config = config_registry.qwen3_235b_a22b() + + assert config.model_spec.flavor == "235B-A22B" + assert config.hf_assets_path == "./assets/hf/Qwen3-235B-A22B" + assert config.training.seq_len == 4096 + assert isinstance(config.activation_checkpoint, FullAC.Config) + + +def test_scaled_235b_model_flavor_is_neuron_owned(): + assert "235B-A22B-scaled" not in qwen3_configs + + config = model_registry("235B-A22B-scaled", attn_backend="sdpa").model + + assert config.dim == 1024 + assert len(config.layers) == 2 + assert config.layers[0].moe.num_experts == 16 + assert config.layers[0].moe.router.top_k == 4 + + +def test_eager_and_compile_recipes_select_the_same_nki_model(): + eager = config_registry.qwen3_30b_a3b_nki_fullac_eager() + compiled = config_registry.qwen3_30b_a3b_nki_fullac_compile() + + for config in (eager, compiled): + assert config.model_spec.flavor == "30B-A3B-nki" + assert len(config.model_spec.model.layers) == 48 + assert config.training.steps == 10 + assert config.dataloader.dataset == "c4_test" + + assert eager.compile.enable is False + assert eager.override.imports == [] + assert compiled.compile.enable is True + assert compiled.compile.backend == "neuron" + assert compiled.compile.components == ["model", "loss", "grad_clip"] + assert compiled.override.imports == [] + + +@pytest.mark.parametrize( + "config_factory", + [ + case + for case in _config_factories() + if "nki" not in case.id and "235b" not in case.id + ], +) +def test_dense_recipes_select_core_prefix_rope(recipe_config): + config = recipe_config + + for layer in config.model_spec.model.layers: + assert type(layer.attention.rope) is CosSinRoPE.Config + assert layer.attention.rope.use_prefix_cache_for_positions is True + assert config.override.imports == [] + if config.compile.enable: + assert "grad_clip" in config.compile.components + + +def test_neuron_recipe_does_not_change_core_model_positions(): + config_registry.qwen3_8b_tp4_fsdp16_sac() + model = core_model_registry("8B").model + + assert all( + not layer.attention.rope.use_prefix_cache_for_positions + for layer in model.layers + ) + + +def _lr(config): + return config.optimizer.param_groups[0].optimizer_kwargs["lr"] + + +def _assert_chunked_cross_entropy_loss(config): + assert isinstance(config.loss, ChunkedLossWrapper.Config) + assert isinstance(config.loss.loss_fn, CrossEntropyLoss.Config) + + +def _assert_sdpa(config): + assert isinstance( + config.model_spec.model.layers[0].attention.inner_attention, + ScaledDotProductAttention.Config, + ) + + +def test_qwen3_8b_dense_eager_contract(): + config = config_registry.qwen3_8b_tp1_fsdp64_fullac_eager() + + assert config.model_spec.name == "qwen3" + assert config.model_spec.flavor == "8B" + _assert_chunked_cross_entropy_loss(config) + _assert_sdpa(config) + assert config.parallelism.tensor_parallel_degree == 1 + assert config.parallelism.data_parallel_shard_degree == -1 + assert config.training.local_batch_size == 1 + assert config.training.global_batch_size == 64 + assert _lr(config) == 8e-4 + assert isinstance(config.activation_checkpoint, FullAC.Config) + assert config.compile.enable is False + assert config.override.imports == [] + + +def test_qwen3_8b_dense_compile_contract(): + config = config_registry.qwen3_8b_tp1_fsdp64_fullac_compile() + + assert config.model_spec.name == "qwen3" + assert config.model_spec.flavor == "8B" + _assert_chunked_cross_entropy_loss(config) + _assert_sdpa(config) + assert config.parallelism.tensor_parallel_degree == 1 + assert config.parallelism.data_parallel_shard_degree == -1 + assert config.training.local_batch_size == 1 + assert config.training.global_batch_size == 64 + assert _lr(config) == 8e-4 + assert isinstance(config.activation_checkpoint, FullAC.Config) + assert config.compile.enable is True + assert config.compile.backend == "neuron" + assert config.compile.components == ["model", "loss", "grad_clip"] + assert config.override.imports == [] + + +def test_qwen3_32b_dense_compile_contract(): + config = config_registry.qwen3_32b_tp4_fsdp16_fullac_compile() + + assert config.model_spec.name == "qwen3" + assert config.model_spec.flavor == "32B" + _assert_chunked_cross_entropy_loss(config) + _assert_sdpa(config) + assert config.parallelism.tensor_parallel_degree == 4 + assert config.parallelism.data_parallel_shard_degree == -1 + assert config.training.local_batch_size == 1 + assert config.training.global_batch_size == 16 + assert _lr(config) == 8e-4 + assert isinstance(config.activation_checkpoint, FullAC.Config) + assert config.compile.enable is True + assert config.compile.backend == "neuron" + assert config.compile.components == ["model", "loss", "grad_clip"] + assert config.override.imports == [] + + +def test_qwen3_235b_production_recipe(): + config = config_registry.qwen3_235b_a22b_neuron() + + assert config.model_spec.flavor == "235B-A22B-neuron" + assert len(config.model_spec.model.layers) == 94 + assert config.training.seq_len == 4096 + assert config.parallelism.tensor_parallel_degree == 4 + assert config.parallelism.data_parallel_shard_degree == 32 + assert config.parallelism.data_parallel_replicate_degree == 4 + assert config.parallelism.pipeline_parallel_degree == 2 + assert config.parallelism.expert_parallel_degree == 128 + assert config.compile.enable is True + assert config.compile.backend == "neuron" + for layer in config.model_spec.model.layers: + assert isinstance(layer.moe.experts, nki_moe.BWMMGroupedExperts.Config) diff --git a/torchtitan_neuron/tests/moe_block_microbench.py b/torchtitan_neuron/tests/moe_block_microbench.py new file mode 100644 index 00000000000..09c2a1df3f5 --- /dev/null +++ b/torchtitan_neuron/tests/moe_block_microbench.py @@ -0,0 +1,463 @@ +#!/usr/bin/env python3 +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Full Qwen3-235B-A22B MoE expert-block microbenchmark (router INCLUDED). + +Unlike moe_dispatcher_microbench.py -- which feeds a pre-synthesized routing +decision straight into dispatch -- this bench runs the WHOLE MoE expert block end +to end, starting from the hidden activation: + + router (gate -> softmax -> top-k -> route_norm) + -> dispatch (router-reorder + permute_a2av) + -> blockwise experts (moe_cte NKI kernel) + -> combine (unpermute_a2av) + +The router is the REAL Qwen3 ``TokenChoiceTopKRouter`` built from the exact +235B-A22B recipe (dim=4096, num_experts=128, top_k=8, score_func="softmax", +route_norm=True); the gate is a live nn.Linear, so expert assignments are +data-dependent on the input activation rather than pre-baked. + +Fidelity note -- there is NO single importable Qwen3-235B block that is +simultaneously (a2av-dispatched) AND (blockwise-NKI experts): + * torchtitan/models/common/moe.py::MoE -> real router + permute/unpermute + a2av, but grouped_mm experts (NOT blockwise). + * torchtitan_neuron/models/qwen3/nki_moe.py::NKIMoE -> real router + blockwise nki_fwd, + but TP-ONLY (raises NotImplementedError under EP; no a2av at all). +So the a2av+blockwise combination only exists as a COMPOSITION. This bench imports +the real Qwen3 router and the real neuron_ep dispatcher (permute_a2av / +unpermute_a2av) and composes them with the blockwise moe_cte experts -- the same +composition validated in moe_dispatcher_microbench.py, now fronted by the router. + +Parallelism (token-count-only fidelity): + Real 235B run = TP4 x DP16 = 64 ranks, EP=64 overlaying all ranks. Under EP, + torchtitan requires sequence-parallel activations, so each rank's router input is + the SP shard of its DP replica's sequence: T_local = seq_per_replica // TP. The + router gate is REPLICATED (not TP-sharded) in the EP path, and expert weights are + pure-EP (this benchmark uses 64 experts, one per rank at EP=64, with no + expert-TP). So TP4's + ONLY effect here is the per-rank token count -- no weight/DTensor sharding is + materialized; each rank builds its own num_local_experts of full expert weights. + + Defaults: seq_per_replica=4096, TP=4 -> T_local=1024 tokens/rank; EP=64 -> + num_local_experts=1; N = T_local*top_k = 8192; packed expert capacity is + align(T_local*EP) = 65536 rows/rank. + +Timing: the full block is ONE timed region with NO intra-block host barriers -- a +single torch.neuron.synchronize() at the end lands router+dispatch+experts+combine +before the timer stops, preserving device overlap. The correctness checksum is taken +outside the timed region (last collected iter only) so its D2H copy never pollutes +timing or a --profile capture. + +Launch (64 ranks == 64 NeuronCores): + torchrun --nproc_per_node=64 moe_block_microbench.py + +Environment variables (shared with moe_dispatcher_microbench.py): + NEURON_RT_VISIBLE_CORES=0-63 + NEURON_USE_NKI_ALLTOALLV=1 + NEURON_LOGICAL_NC_CONFIG=2 + TORCH_NEURONX_ENABLE_LAZY_ALLOC=0 + NEURON_RT_EXEC_TIMEOUT=120 +""" + +from __future__ import annotations + +import argparse +import os +import statistics +import sys +import time +from functools import partial + +import torch +import torch.distributed as dist + +# Registers the "neuron" device + distributed backend + the mesh_registry hook +# that neuron_ep's build_ep_replica_group_ranks reads. +import torch_neuronx # noqa: F401 +from torch import nn +from torch.distributed.device_mesh import init_device_mesh + + +# ---- Qwen3-235B-A22B MoE recipe (torchtitan/models/qwen3/__init__.py::_235b_a22b) ---- +H = 4096 # dim (D) +FFN = 1536 # moe_hidden_dim (F); per-expert intermediate +TOP_K = 8 +NUM_EXPERTS = 64 # reduced from the 235B recipe's 128 (at EP=64 -> num_local_experts=1) +SCORE_FUNC = "softmax" +ROUTE_NORM = True +DEFAULT_SEQ_PER_REPLICA = 4096 # sequence length per DP replica +DEFAULT_TP = 4 # tensor-parallel degree (folds into per-rank token count) +DEFAULT_EP = 64 # expert-parallel group size (= world for pure EP) +DEFAULT_BLOCK_SIZE = 512 + + +def log(rank, msg): + print(f"[rank {rank}] {msg}", flush=True) + sys.stdout.flush() + + +def build_router(device): + """Build the REAL Qwen3-235B TokenChoiceTopKRouter and initialize the gate. + + Uses the same factory the model uses (make_router_config) with the 235B recipe. + The gate is a live nn.Linear(H, NUM_EXPERTS); trunc_normal init mirrors the + model's _LINEAR_INIT so scores are realistic. Router runs on the (replicated) EP + path, so no TP sharding is applied. + """ + from torchtitan.models.common.config_utils import make_router_config + + router_cfg = make_router_config( + dim=H, + num_experts=NUM_EXPERTS, + gate_param_init={ + "weight": partial(nn.init.trunc_normal_, std=0.02), + "bias": nn.init.zeros_, + }, + top_k=TOP_K, + score_func=SCORE_FUNC, + route_norm=ROUTE_NORM, + ) + router = router_cfg.build() + router.init_states(buffer_device=torch.device(device)) + return router.to(device) + + +def route(router, x_TD): + """Run the real router and derive the dispatcher's routing-decision contract. + + Mirrors MoE.forward's router glue (torchtitan/models/common/moe.py): the router + returns per-token top-k scores and expert ids; num_local_tokens_per_expert_E is + the per-expert routed-token count from the one-hot routing map. Returns + (topk_scores_TK f32, topk_expert_ids_TK i32, num_local_tokens_per_expert_E i32). + """ + T, D = x_TD.shape + # Router expects (B, L, D); treat the rank's tokens as one (1, T, D) batch. + x_BLD = x_TD.view(1, T, D) + topk_scores_BLK, topk_expert_ids_BLK, scores_BLE = router(x_BLD, None) + + topk_scores_TK = topk_scores_BLK.view(T, TOP_K).to(torch.float32) + topk_expert_ids_TK = topk_expert_ids_BLK.view(T, TOP_K).to(torch.int32) + + routing_map_TE = torch.zeros_like(scores_BLE, dtype=torch.bool).view(T, NUM_EXPERTS) + routing_map_TE.scatter_(-1, topk_expert_ids_TK.to(torch.int64), True) + num_local_tokens_per_expert_E = routing_map_TE.sum(dim=0).to(torch.int32) + return topk_scores_TK, topk_expert_ids_TK, num_local_tokens_per_expert_E + + +def make_expert_weights(num_local_experts, device): + """Random per-expert SwiGLU weights in GroupedExperts param shapes (bf16-ready).""" + g = torch.Generator().manual_seed(42) + scale = 1.0 / (H**0.5) + w1_EFD = (torch.randn(num_local_experts, FFN, H, generator=g) * scale).to(device) + w3_EFD = (torch.randn(num_local_experts, FFN, H, generator=g) * scale).to(device) + w2_EDF = (torch.randn(num_local_experts, H, FFN, generator=g) * scale).to(device) + return w1_EFD, w3_EFD, w2_EDF + + +def repack_weights_for_moe_cte(weights): + """Repack grouped-mm expert weights into moe_cte's layout. + + Same as moe_dispatcher_microbench.repack_weights_for_moe_cte: + w1/w3_EFD [E,FFN,H], w2_EDF [E,H,FFN] -> + gate_up [E, H, 2, FFN] (0=gate 1=up), down [E, FFN, H], both bf16. + """ + w1_EFD, w3_EFD, w2_EDF = weights + gate = w1_EFD.transpose(-2, -1) # [E, H, FFN] + up = w3_EFD.transpose(-2, -1) # [E, H, FFN] + gate_up = ( + torch.stack([gate, up], dim=2).contiguous().to(torch.bfloat16) + ) # [E,H,2,FFN] + down = w2_EDF.transpose(-2, -1).contiguous().to(torch.bfloat16) # [E,FFN,H] + return gate_up, down + + +def moe_cte_experts_forward( + x_RD, num_local_experts, gate_up_w, down_w, block_size, spec, lnc +): + """Blockwise-NKI expert compute over the padded expert-major dispatch buffer. + + Identical to moe_dispatcher_microbench.moe_cte_experts_forward: x_RD is + [num_local_experts * (EP*cap), H] with each expert's rows contiguous, so block + metadata is IDENTITY (block b -> expert (b*B)//rows_per_expert, + token_position_to_id = arange(R), affinity = 1; routing scores applied in + combine). moe_cte tiles internally by block_size, so it compiles where a single + grouped-mm over R rows would overflow the neuronx-cc 5M-instruction limit. + """ + from nki_blockwise_moe_cte import wrapped_moe_cte # shared Neuron test helper + + R, Hd = x_RD.shape + device = x_RD.device + B = block_size + assert R % B == 0, f"R={R} not divisible by block_size={B}" + rows_per_expert = R // num_local_experts + assert rows_per_expert % B == 0, "rows_per_expert must be a block_size multiple" + n_blk = R // B + + hidden = torch.cat([x_RD, x_RD.new_zeros(1, Hd)], dim=0) # [R+1, H] (dummy pad row) + aff = torch.zeros(R + 1, num_local_experts, dtype=torch.float32, device=device) + row_expert = torch.arange(R, device=device) // rows_per_expert + aff[torch.arange(R, device=device), row_expert] = 1.0 + aff = aff.reshape((R + 1) * num_local_experts, 1) + tpid = torch.arange(R, device=device, dtype=torch.int32) # identity, all rows real + blk_expert = torch.arange(n_blk, device=device) * B // rows_per_expert + block_to_expert = blk_expert.to(torch.int32).reshape(-1, 1) + + out = wrapped_moe_cte[(lnc,)]( + hidden, + aff, + gate_up_w, + down_w, + tpid, + block_to_expert, + B, + spec, + ) + out = out[0] if isinstance(out, tuple) else out + return out[:R].type_as(x_RD) + + +def summarize(name, samples_s, rank): + """Log mean / p50 / p90 for a phase (samples in seconds -> ms).""" + ms = sorted(s * 1e3 for s in samples_s) + mean = statistics.mean(ms) + p50 = ms[len(ms) // 2] + p90 = ms[min(len(ms) - 1, int(0.9 * len(ms)))] + log( + rank, + f" {name:<12} mean={mean:8.2f}ms p50={p50:8.2f}ms p90={p90:8.2f}ms (n={len(ms)})", + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--ep", type=int, default=DEFAULT_EP, help="expert-parallel group size" + ) + ap.add_argument( + "--tp", + type=int, + default=DEFAULT_TP, + help="tensor-parallel degree; only sets per-rank token count " + "(T_local = seq-per-replica // tp)", + ) + ap.add_argument( + "--seq-per-replica", + type=int, + default=DEFAULT_SEQ_PER_REPLICA, + help="sequence length per DP replica (before SP sharding)", + ) + ap.add_argument( + "--block-size", type=int, default=DEFAULT_BLOCK_SIZE, help="moe_cte block size" + ) + ap.add_argument("--warmup", type=int, default=3) + ap.add_argument("--iters", type=int, default=10) + ap.add_argument( + "--profile", + action="store_true", + help="capture Neuron DEVICE (.ntff) + RUNTIME (system) traces after timing", + ) + ap.add_argument( + "--profile-dir", + default=None, + help="profile output dir (default: $MB_PROFILE_DIR or ./mb_profile)", + ) + ap.add_argument( + "--profile-active", + type=int, + default=2, + help="profiled steady-state steps to capture", + ) + args = ap.parse_args() + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + ep = args.ep + dist.init_process_group(backend="neuron") + + assert world_size % ep == 0, f"world_size {world_size} not divisible by ep {ep}" + assert NUM_EXPERTS % ep == 0, f"num_experts {NUM_EXPERTS} not divisible by ep {ep}" + assert ep == NUM_EXPERTS, ( + f"packed neuron_ep requires one local expert per rank: " + f"ep ({ep}) must equal num_experts ({NUM_EXPERTS})." + ) + assert ( + args.seq_per_replica % args.tp == 0 + ), f"seq-per-replica {args.seq_per_replica} not divisible by tp {args.tp}" + num_local_experts = NUM_EXPERTS // ep + # SP shards the DP-replica sequence across the TP group: each rank routes T_local. + T = args.seq_per_replica // args.tp + device = "neuron" + + # Real neuron_ep dispatcher (permute_a2av / unpermute_a2av), packed capacity. + # Build from the package factory so the benchmark exercises the package implementation + # directly, with no standalone copy to drift. + from torchtitan.models.common.config_utils import make_token_dispatcher_config + + dispatcher = make_token_dispatcher_config( + num_experts=NUM_EXPERTS, top_k=TOP_K, comm_backend="neuron_ep" + ).build() + + # 2D (efsdp, ep) mesh; ep innermost so replica groups are contiguous, matching the + # parallel_dims sparse-mesh ordering. wire_meshes(tp_mesh=None): experts are pure-EP. + mesh = init_device_mesh( + "neuron", (world_size // ep, ep), mesh_dim_names=("efsdp", "ep") + ) + ep_mesh = mesh["ep"] + dispatcher.wire_meshes(ep_mesh=ep_mesh, tp_mesh=None) + + from torch_neuronx.utils import get_logical_neuron_cores + + lnc = int(get_logical_neuron_cores()) + + router = build_router(device) + x_TD = torch.randn(T, H, dtype=torch.bfloat16, device=device) + weights = make_expert_weights(num_local_experts, device) + moe_cte_gw, moe_cte_dw = repack_weights_for_moe_cte(weights) + import nki_blockwise_moe_cte as _bw + + moe_cte_spec = _bw.default_spec() + + if rank == 0: + log( + rank, + f"Qwen3-235B-A22B MoE block world={world_size} ep={ep} tp={args.tp} " + f"local_experts={num_local_experts} seq_per_replica={args.seq_per_replica} " + f"T_local={T} H={H} FFN={FFN} top_k={TOP_K} score={SCORE_FUNC} " + f"route_norm={ROUTE_NORM}", + ) + log( + rank, + f"experts=moe_cte block_size={args.block_size} lnc={lnc} " + f"warmup={args.warmup} iters={args.iters}", + ) + + def one_step(collect): + """Run the FULL MoE block once: router -> dispatch -> blockwise experts -> + combine. No intra-block host barriers; a single torch.neuron.synchronize() + lands the whole pipeline before the timer stops. Checksum is taken outside + the timed region (only when ``collect``) so its D2H never pollutes timing. + """ + t0 = time.time() + # Router: gate -> softmax -> top-k -> route_norm (data-dependent on x). + scores_TK, ids_TK, counts_E = route(router, x_TD) + # Dispatch: router-reorder + permute_a2av. + routed_input_RD, _counts_e, meta = dispatcher.dispatch( + x_TD, scores_TK, ids_TK, counts_E + ) + # Experts: blockwise moe_cte NKI kernel over the padded expert-major buffer. + routed_output_RD = moe_cte_experts_forward( + routed_input_RD, + num_local_experts, + moe_cte_gw, + moe_cte_dw, + args.block_size, + moe_cte_spec, + lnc, + ) + # Combine: unpermute_a2av (+ scatter_add with routing scores). + out_TD = dispatcher.combine( + routed_output_RD, + meta, + x_TD, + num_local_tokens_after_padding=T, + local_seq_len_after_padding=T, + ) + torch.neuron.synchronize() # single barrier for the whole block + t1 = time.time() + if collect: + checksum = float(out_TD.detach().to("cpu").float().abs().sum().item()) + return (t1 - t0,), checksum + return None, None + + # ---- Warmup (compile + cache the NEFFs; excluded from timing) ---- + dist.barrier() + log(rank, f"=== warmup x{args.warmup} ===") + for _ in range(args.warmup): + one_step(collect=False) + dist.barrier() + + # ---- Timed iters (full MoE-block latency) ---- + log(rank, f"=== timing x{args.iters} ===") + full = [] + last_checksum = None + for _ in range(args.iters): + times, last_checksum = one_step(collect=True) + full.append(times[0]) + dist.barrier() + + log( + rank, + f"=== RESULTS Qwen3-235B MoE block ep={ep} checksum={last_checksum:.3e} ===", + ) + summarize("full-block", full, rank) + + if args.profile: + _capture_profile(args, rank, one_step) + + dist.barrier() + if rank == 0: + log(rank, "ALL RANKS COMPLETED SUCCESSFULLY (Qwen3-235B MoE block)") + dist.destroy_process_group() + + +def _capture_profile(args, rank, one_step): + """Capture DEVICE (.ntff hardware) + RUNTIME (system) traces for Neuron Explorer. + + Same knobs as moe_dispatcher_microbench._capture_profile: both DEVICE and RUNTIME + modes plus CPU_UTIL/HOST_MEMORY, oversized event ring buffer to avoid silent drops, + and both LNC cores captured. NEFFs are already warm from the timing loop. + """ + from torch.profiler import profile, ProfilerActivity, schedule + from torch_neuronx.profiling import NeuronConfig, NeuronProfiler, ProfileMode + + profile_dir = args.profile_dir or os.environ.get("MB_PROFILE_DIR") or "./mb_profile" + rank_dir = os.path.join( + profile_dir, f"qwen3_235b_moe_block_ep{args.ep}", f"rank{rank}" + ) + os.makedirs(rank_dir, exist_ok=True) + neuron_config = NeuronConfig( + modes=[ + ProfileMode.DEVICE, + ProfileMode.RUNTIME, + ProfileMode.CPU_UTIL, + ProfileMode.HOST_MEMORY, + ], + profile_output_dir=rank_dir, + max_events_per_nc=10_000_000, + capture_enabled_for_nc="0,1", + ) + exporter = NeuronProfiler(neuron_config) + + if rank == 0: + log( + rank, + f"=== PROFILING (DEVICE+RUNTIME) active={args.profile_active} dir={profile_dir} ===", + ) + dist.barrier() + sched = schedule(wait=1, warmup=1, active=args.profile_active, repeat=1) + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1], + experimental_config=neuron_config, + schedule=sched, + on_trace_ready=exporter.export_trace, + record_shapes=True, + ) as prof: + for _ in range(2 + args.profile_active): + one_step(collect=False) + torch.neuron.synchronize() + prof.step() + try: + prof.export_chrome_trace(os.path.join(rank_dir, f"framework_rank{rank}.json")) + except Exception as e: # noqa: BLE001 + log(rank, f"chrome-trace export skipped: {e}") + dist.barrier() + if rank == 0: + log(rank, f"=== PROFILE session dir: {neuron_config.get_session_dir()} ===") + log(rank, "=== PROFILE DONE ===") + + +if __name__ == "__main__": + main() diff --git a/torchtitan_neuron/tests/moe_dispatcher_microbench.py b/torchtitan_neuron/tests/moe_dispatcher_microbench.py new file mode 100644 index 00000000000..ccc73670833 --- /dev/null +++ b/torchtitan_neuron/tests/moe_dispatcher_microbench.py @@ -0,0 +1,567 @@ +#!/usr/bin/env python3 +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Dispatcher-agnostic 64-rank MoE token-dispatch microbenchmark. + +Runs the full MoE expert path in isolation -- NO model, NO router network: + + dispatch(tokens) -> grouped-mm (SwiGLU expert compute) -> combine(outputs) + +against ANY Neuron-supported token dispatcher, selected by --backend: + + * standard -> AllToAllTokenDispatcher (c10d all_to_all_single, NKI a2a-v under + NEURON_USE_NKI_ALLTOALLV=1) + * neuron_ep -> NeuronEPTokenDispatcher (NKI permute_a2av / unpermute_a2av) + * local -> LocalTokenDispatcher (EP=1 baseline, no collective) + +All three implement the SAME dispatch()/combine() contract from +torchtitan/models/common/token_dispatcher.py, so the harness body is identical +across backends -- the only difference is the factory line. The grouped-mm in the +middle mirrors GroupedExperts._experts_forward in moe.py (silu(x@w1) * (x@w3) then +@w2, with offs = cumsum(per-expert counts)), including the Neuron 128-alignment +padding that the standard backend needs (neuron_ep pre-pads in its dispatcher). + +It measures the FULL expert block as a single timed region -- + + dispatch(router-reorder + permute_a2av) -> blockwise experts -> combine(unpermute_a2av) + +-- with NO intra-block host barriers: one torch.neuron.synchronize() at the end of the +block forces the whole device pipeline to land before the timer is stamped. This is a +true full-block latency (device dispatch + collectives + compute overlap intact), unlike +per-phase timing which needs a D2H sync after every phase and serializes host/device. +Reports mean / p50 / p90 per rank. A correctness checksum is taken on the LAST iter only, +AFTER the timer stops, so its D2H copy never pollutes timing or a --profile capture. + +Launch (64 ranks == 64 NeuronCores): + torchrun --nproc_per_node=64 moe_dispatcher_microbench.py --backend neuron_ep + torchrun --nproc_per_node=64 moe_dispatcher_microbench.py --backend standard + torchrun --nproc_per_node=64 moe_dispatcher_microbench.py --backend local --ep 1 + +Env (set by pod_run_moe_dispatcher_microbench.sh): + NEURON_RT_VISIBLE_CORES=0-63 + NEURON_USE_NKI_ALLTOALLV=1 + NEURON_LOGICAL_NC_CONFIG=2 + TORCH_NEURONX_ENABLE_LAZY_ALLOC=0 + NEURON_RT_EXEC_TIMEOUT=120 +""" + +from __future__ import annotations + +import argparse +import os +import statistics +import sys +import time + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +# Registers the "neuron" device + distributed backend + the mesh_registry hook +# that neuron_ep's build_ep_replica_group_ranks reads. +import torch_neuronx # noqa: F401 +from torch.distributed.device_mesh import init_device_mesh + + +# ---- Model-shaped defaults (Qwen3-235B-A22B, matches the a2av microbenches) ---- +H = 4096 # hidden dim (D) +FFN = 1536 # per-expert intermediate dim (F); Qwen3-235B moe_inter_dim +TOP_K = 8 +NUM_EXPERTS = 64 # reduced from the 235B recipe's 128 (at EP=64 -> num_local_experts=1) +DEFAULT_SEQ_LEN = 512 # T tokens / rank +DEFAULT_EP = 64 # packed neuron_ep requires EP >= 8 and one local expert + + +def log(rank, msg): + print(f"[rank {rank}] {msg}", flush=True) + sys.stdout.flush() + + +def make_dispatcher(backend: str): + """Return an unbuilt dispatcher for the requested Neuron-supported backend. + + Uses the SAME factory the model uses (config_utils.make_token_dispatcher_config), + so this stays in lockstep with whatever backends torchtitan supports. ``local`` + is an alias for ``standard`` at EP=1 -- the AllToAllTokenDispatcher falls back to + LocalTokenDispatcher when ``ep_mesh is None`` (see AllToAllTokenDispatcher.dispatch), + so passing ``standard`` with ep=1 exercises the exact local-dispatch code path. + """ + # neuron_ep packed A2AV is a first-class backend in the torchtitan factory, + # so it flows through the generic path below like every other backend. + from torchtitan.models.common.config_utils import make_token_dispatcher_config + + factory_backend = "standard" if backend == "local" else backend + cfg = make_token_dispatcher_config( + num_experts=NUM_EXPERTS, top_k=TOP_K, comm_backend=factory_backend + ) + return cfg.build() + + +def build_router(rank: int, seq_len: int, imbalance: str, device): + """Synthesize a router decision (no gate network): per-token top_k expert ids. + + imbalance: + uniform -> every rank's tokens spread evenly over all experts + skewed -> per-rank random routing (data-dependent counts, mimics a real + router's load imbalance across EP ranks) + Returns (topk_scores_TK float32, topk_expert_ids_TK int32, + num_local_tokens_per_expert_E int32) on ``device``. + """ + g = torch.Generator().manual_seed(3000 + rank) + if imbalance == "uniform": + # Round-robin: exact balance across experts. + ids_TK = ( + torch.arange(seq_len * TOP_K, dtype=torch.int64) % NUM_EXPERTS + ).reshape(seq_len, TOP_K) + else: + # Each token picks top_k DISTINCT experts; per-rank seed -> imbalance. + ids_TK = torch.stack( + [torch.randperm(NUM_EXPERTS, generator=g)[:TOP_K] for _ in range(seq_len)] + ) + ids_TK = ids_TK.to(torch.int32) + scores_TK = torch.rand(seq_len, TOP_K, generator=g, dtype=torch.float32) + counts_E = torch.bincount( + ids_TK.reshape(-1).to(torch.int64), minlength=NUM_EXPERTS + ).to(torch.int32) + return ( + scores_TK.to(device), + ids_TK.to(device), + counts_E.to(device), + ) + + +def make_expert_weights(num_local_experts: int, device): + """Random per-expert SwiGLU weights, matching GroupedExperts param shapes.""" + g = torch.Generator().manual_seed(42) + scale = 1.0 / (H**0.5) + w1_EFD = (torch.randn(num_local_experts, FFN, H, generator=g) * scale).to(device) + w3_EFD = (torch.randn(num_local_experts, FFN, H, generator=g) * scale).to(device) + w2_EDF = (torch.randn(num_local_experts, H, FFN, generator=g) * scale).to(device) + return w1_EFD, w3_EFD, w2_EDF + + +def gmm_align_value() -> int: + from torchtitan.distributed.neuron_ep import gmm_align + + return gmm_align() + + +def pad_experts_for_neuron(x_RD, counts_e, align): + """Pad each expert's token block to an ``align`` multiple for this benchmark. + + grouped_mm on Neuron requires the token dim and every offs boundary to be + align-divisible. neuron_ep already pre-pads (counts come back aligned, so this + is a no-op); the standard/local backends return raw counts and need it here. + + Returns (x_padded, padded_counts_e, gather_R) where gather_R selects the R real + rows back out of the padded output (None when no padding was needed). + """ + real_counts_e = counts_e.to(torch.int64) + padded_counts_e = torch.ceil(real_counts_e / align).to(torch.int64) * align + if bool((padded_counts_e == real_counts_e).all()): + return x_RD, counts_e, None + + device = x_RD.device + num_experts = real_counts_e.shape[0] + r_pad = int(padded_counts_e.sum()) + real_starts_e = real_counts_e.cumsum(0) - real_counts_e + padded_starts_e = padded_counts_e.cumsum(0) - padded_counts_e + expert_of_row_R = torch.arange(num_experts, device=device).repeat_interleave( + real_counts_e + ) + padded_positions_R = ( + padded_starts_e[expert_of_row_R] + + torch.arange(expert_of_row_R.shape[0], device=device) + - real_starts_e[expert_of_row_R] + ) + x_padded = x_RD.new_zeros((r_pad, x_RD.shape[-1])) + x_padded[padded_positions_R] = x_RD + return x_padded, padded_counts_e.to(torch.int32), padded_positions_R + + +def repack_weights_for_moe_cte(weights): + """Repack grouped-mm expert weights into moe_cte's layout. + + make_expert_weights gives w1_EFD/w3_EFD [E,FFN,H] and w2_EDF [E,H,FFN] (the + grouped-mm convention: weight applied as x @ w.transpose(-2,-1)). moe_cte wants: + gate_up_proj_weight [E, H, 2, I_TP] (h @ gate_up[e] -> [.,2*I_TP]; 0=gate 1=up) + down_proj_weight [E, I_TP, H] (inter @ down[e] -> [.,H]) + so gate/up = w1/w3 transposed to [H,FFN]; down = w2 transposed to [FFN,H]. + """ + w1_EFD, w3_EFD, w2_EDF = weights + gate = w1_EFD.transpose(-2, -1) # [E, H, FFN] + up = w3_EFD.transpose(-2, -1) # [E, H, FFN] + gate_up = ( + torch.stack([gate, up], dim=2).contiguous().to(torch.bfloat16) + ) # [E,H,2,FFN] + down = w2_EDF.transpose(-2, -1).contiguous().to(torch.bfloat16) # [E,FFN,H] + return gate_up, down + + +def moe_cte_experts_forward( + x_RD, num_local_experts, gate_up_w, down_w, block_size, spec, lnc +): + """Blockwise-NKI expert compute over the padded expert-major dispatch buffer. + + Replaces the grouped-mm ``experts_forward`` for neuron_ep: ``x_RD`` is + ``[num_local_experts * rows_per_expert, H]`` with each expert's rows contiguous + (rows_per_expert = EP*cap). Since rows are already expert-grouped and R is a + block_size multiple, block metadata is IDENTITY: block b -> expert (b*B)//rpe, + token_position_to_id = arange(R). Affinity = 1 (routing scores are applied later in + combine, matching the grouped-mm path). Runs moe_cte (tiles internally by block_size), + so it COMPILES where the single grouped-mm over R rows overflowed the 5M limit. + """ + from nki_blockwise_moe_cte import wrapped_moe_cte # shared Neuron test helper + + R, Hd = x_RD.shape + device = x_RD.device + B = block_size + assert R % B == 0, f"R={R} not divisible by block_size={B}" + rows_per_expert = R // num_local_experts + assert rows_per_expert % B == 0, "rows_per_expert must be a block_size multiple" + n_blk = R // B + + hidden = torch.cat([x_RD, x_RD.new_zeros(1, Hd)], dim=0) # [R+1, H] (dummy pad row) + # affinity[row, expert_of_row] = 1; padding row 0. Flattened [(R+1)*e, 1]. + aff = torch.zeros(R + 1, num_local_experts, dtype=torch.float32, device=device) + row_expert = torch.arange(R, device=device) // rows_per_expert + aff[torch.arange(R, device=device), row_expert] = 1.0 + aff = aff.reshape((R + 1) * num_local_experts, 1) + tpid = torch.arange(R, device=device, dtype=torch.int32) # identity, all rows real + blk_expert = torch.arange(n_blk, device=device) * B // rows_per_expert + block_to_expert = blk_expert.to(torch.int32).reshape(-1, 1) + + out = wrapped_moe_cte[(lnc,)]( + hidden, + aff, + gate_up_w, + down_w, + tpid, + block_to_expert, + B, + spec, + ) + out = out[0] if isinstance(out, tuple) else out + return out[:R].type_as(x_RD) + + +def experts_forward(x_RD, counts_e, weights, align): + """SwiGLU grouped-mm expert compute -- mirror of GroupedExperts._experts_forward. + + dispatch gives expert-major rows + per-local-expert counts; run the 3 grouped + matmuls with offs = cumsum(counts), then drop any Neuron alignment padding so + the (R, D) -> (R, D) contract combine expects is preserved. + """ + w1_EFD, w3_EFD, w2_EDF = weights + x_RD, counts_e, gather_R = pad_experts_for_neuron(x_RD, counts_e, align) + offsets_E = torch.cumsum(counts_e, dim=0, dtype=torch.int32) + + h_RF = F.silu( + torch._grouped_mm( + x_RD.bfloat16(), w1_EFD.bfloat16().transpose(-2, -1), offs=offsets_E + ) + ) + h_RF = h_RF * torch._grouped_mm( + x_RD.bfloat16(), w3_EFD.bfloat16().transpose(-2, -1), offs=offsets_E + ) + out_RD = torch._grouped_mm( + h_RF, w2_EDF.bfloat16().transpose(-2, -1), offs=offsets_E + ).type_as(x_RD) + + if gather_R is not None: + out_RD = out_RD[gather_R] + return out_RD + + +def summarize(name, samples_s, rank): + """Log mean / p50 / p90 for a phase (samples in seconds -> ms).""" + ms = sorted(s * 1e3 for s in samples_s) + mean = statistics.mean(ms) + p50 = ms[len(ms) // 2] + p90 = ms[min(len(ms) - 1, int(0.9 * len(ms)))] + log( + rank, + " {name:<12} mean={mean:8.2f}ms p50={p50:8.2f}ms " + "p90={p90:8.2f}ms (n={count})".format( + name=name, + mean=mean, + p50=p50, + p90=p90, + count=len(ms), + ), + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--backend", choices=["standard", "neuron_ep", "local"], default="neuron_ep" + ) + ap.add_argument( + "--ep", type=int, default=DEFAULT_EP, help="expert-parallel group size" + ) + ap.add_argument( + "--seq-len", type=int, default=DEFAULT_SEQ_LEN, help="T tokens per rank" + ) + ap.add_argument("--imbalance", choices=["uniform", "skewed"], default="skewed") + ap.add_argument( + "--experts", + choices=["grouped_mm", "moe_cte"], + default="grouped_mm", + help="expert compute: grouped-mm (default) or blockwise moe_cte NKI kernel", + ) + ap.add_argument("--block-size", type=int, default=512, help="moe_cte block size") + ap.add_argument("--warmup", type=int, default=3) + ap.add_argument("--iters", type=int, default=10) + ap.add_argument( + "--profile", + action="store_true", + help="capture Neuron DEVICE (.ntff hardware) + RUNTIME (system) " + "traces after timing, for Neuron Explorer", + ) + ap.add_argument( + "--profile-dir", + default=None, + help="profile output dir (default: $MB_PROFILE_DIR or ./mb_profile)", + ) + ap.add_argument( + "--profile-active", + type=int, + default=2, + help="profiled steady-state steps to capture", + ) + args = ap.parse_args() + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + ep = 1 if args.backend == "local" else args.ep + dist.init_process_group(backend="neuron") + + assert world_size % ep == 0, f"world_size {world_size} not divisible by ep {ep}" + assert NUM_EXPERTS % ep == 0, f"num_experts {NUM_EXPERTS} not divisible by ep {ep}" + num_local_experts = NUM_EXPERTS // ep + device = "neuron" + + dispatcher = make_dispatcher(args.backend) + + # Build a 2D (efsdp, ep) mesh; ep is the innermost axis so replica groups are + # contiguous, exactly matching parallel_dims sparse-mesh ordering. Wire ep_mesh + # into the dispatcher (None for local/EP=1 -> local-dispatch fallback). + ep_mesh = None + if ep > 1: + mesh = init_device_mesh( + "neuron", (world_size // ep, ep), mesh_dim_names=("efsdp", "ep") + ) + ep_mesh = mesh["ep"] + dispatcher.wire_meshes(ep_mesh=ep_mesh, tp_mesh=None) + + if rank == 0: + log( + rank, + f"backend={args.backend} world={world_size} ep={ep} " + f"local_experts={num_local_experts} T={args.seq_len} H={H} FFN={FFN} " + f"top_k={TOP_K} imbalance={args.imbalance}", + ) + log( + rank, + f"warmup={args.warmup} iters={args.iters} gmm_align={gmm_align_value()}", + ) + + scores_TK, ids_TK, counts_E = build_router( + rank, args.seq_len, args.imbalance, device + ) + x_TD = torch.randn(args.seq_len, H, dtype=torch.bfloat16, device=device) + weights = make_expert_weights(num_local_experts, device) + align = gmm_align_value() + T = args.seq_len + + # Expert-compute path: grouped-mm (default) or the blockwise moe_cte NKI kernel. + # moe_cte tiles internally by block_size, so it compiles where the single grouped-mm + # over the neuron_ep padded buffer (num_local_experts*EP*cap rows) + # overflows the neuronx-cc 5M-instruction limit. + moe_cte_gw = moe_cte_dw = moe_cte_spec = None + from torch_neuronx.utils import get_logical_neuron_cores + + lnc = int(get_logical_neuron_cores()) + if args.experts == "moe_cte": + # moe_cte_experts_forward assumes the neuron_ep fixed-capacity layout + # (each local expert = exactly EP*cap contiguous rows, zero-padded). The + # standard/local dispatchers return variable cumsum-packed buffers, so + # rows_per_expert = R // num_local_experts and the block->expert mapping would be + # WRONG (silently, if the block-size asserts happen to pass). Require the matching + # backend rather than risk a wrong answer. + if args.backend != "neuron_ep": + raise ValueError( + f"--experts moe_cte requires --backend neuron_ep (got " + f"{args.backend}). The moe_cte glue assumes the packed " + "one-local-expert capacity layout." + ) + import nki_blockwise_moe_cte as _bw + + moe_cte_gw, moe_cte_dw = repack_weights_for_moe_cte(weights) + moe_cte_spec = _bw.default_spec() + if rank == 0: + log(rank, f"experts=moe_cte block_size={args.block_size} lnc={lnc}") + + def one_step(collect): + """Run the full expert block once: dispatch (router-reorder + permute_a2av) -> + blockwise experts -> combine (unpermute_a2av). + + NO intra-block host barriers -- the phases queue on the device back-to-back and a + single torch.neuron.synchronize() at the end forces the whole pipeline to land + before the timer stops, so device dispatch/collective/compute overlap is preserved. + The correctness checksum is taken outside the timed region (and only when + ``collect`` is set, i.e. never during the timing loop's steady state -- see caller) + so its D2H copy never pollutes timing or a --profile capture. + """ + t0 = time.time() + routed_input_RD, counts_e, meta = dispatcher.dispatch( + x_TD, scores_TK, ids_TK, counts_E + ) + + if args.experts == "moe_cte": + routed_output_RD = moe_cte_experts_forward( + routed_input_RD, + num_local_experts, + moe_cte_gw, + moe_cte_dw, + args.block_size, + moe_cte_spec, + lnc, + ) + else: + routed_output_RD = experts_forward( + routed_input_RD, counts_e, weights, align + ) + + out_TD = dispatcher.combine( + routed_output_RD, + meta, + x_TD, + num_local_tokens_after_padding=T, + local_seq_len_after_padding=T, + ) + # Single barrier for the whole block -- lands dispatch + experts + combine. + torch.neuron.synchronize() + t1 = time.time() + if collect: + checksum = float(out_TD.detach().to("cpu").float().abs().sum().item()) + return (t1 - t0,), checksum + return None, None + + # ---- Warmup (compile + cache the NEFFs; excluded from timing) ---- + dist.barrier() + log(rank, f"=== warmup x{args.warmup} ===") + for _ in range(args.warmup): + one_step(collect=False) + dist.barrier() + + # ---- Timed iters (full expert-block latency) ---- + log(rank, f"=== timing x{args.iters} ===") + full = [] + last_checksum = None + for _ in range(args.iters): + times, last_checksum = one_step(collect=True) + full.append(times[0]) + dist.barrier() + + log( + rank, + "=== RESULTS backend={} ep={} checksum={:.3e} ===".format( + args.backend, ep, last_checksum + ), + ) + summarize("full-block", full, rank) + + # ---- Optional Neuron device + system profile (post-timing, NEFFs already cached) ---- + if args.profile: + _capture_profile(args, rank, one_step) + + dist.barrier() + if rank == 0: + log(rank, f"ALL RANKS COMPLETED SUCCESSFULLY (backend={args.backend})") + dist.destroy_process_group() + + +def _capture_profile(args, rank, one_step): + """Capture DEVICE (.ntff hardware) + RUNTIME (system) traces for Neuron Explorer. + + Uses NeuronProfiler with both ProfileMode.DEVICE and ProfileMode.RUNTIME, driven by + a torch.profiler.schedule so a couple of steady-state steps are captured (NEFFs are + already warm from the timing loop above). Each rank writes its own session under + profile_dir; post-process rank 0 with `neuron-explorer view` into a .pftrace. DEVICE + mode reserves ~5GB HBM per NC -- fine at these sizes. + """ + from torch.profiler import profile, ProfilerActivity, schedule + from torch_neuronx.profiling import NeuronConfig, NeuronProfiler, ProfileMode + + profile_dir = args.profile_dir or os.environ.get("MB_PROFILE_DIR") or "./mb_profile" + # Per-rank subdir so 64 ranks don't clobber one session dir. + rank_dir = os.path.join( + profile_dir, + f"{args.backend}_ep{args.ep if args.backend != 'local' else 1}", + f"rank{rank}", + ) + os.makedirs(rank_dir, exist_ok=True) + # max_events_per_nc sizes the system-trace ring buffer. The NeuronConfig default + # (100K events/NC) silently drops the large majority of events on a 64-rank MoE step + # (~837K events/step measured on a TP4/FSDP16 run -> ~86% dropped at 100K). + # 10M gives headroom. This is the real anti-drop knob -- the + # runtime drops SILENTLY (no log message), so oversize is the only prevention. + # capture_enabled_for_nc="0,1" captures both cores of the LNC=2 group (the guide's + # host-cc example sets it; default may capture NC 0 only). CPU_UTIL/HOST_MEMORY add + # the cpu_util.pb / host_mem.pb host-side context traces. + # DEVICE mode captures the .ntff hardware trace but its notification capture can + # collide with the in-kernel NKI a2a-v collective coordination (permute_a2av does + # cross-rank all_to_all_v INSIDE the kernel) -> ranks desync -> NRT status 1200 + # "Stuck in collectives op (missing notification(s))". Set MB_PROFILE_RUNTIME_ONLY=1 + # to drop DEVICE mode and capture only the system (ntrace.pb) + host traces, which do + # not perturb the collective. RUNTIME still shows dispatch/collective/combine timing. + _runtime_only = os.environ.get("MB_PROFILE_RUNTIME_ONLY") == "1" + _modes = [ProfileMode.RUNTIME, ProfileMode.CPU_UTIL, ProfileMode.HOST_MEMORY] + if not _runtime_only: + _modes = [ProfileMode.DEVICE] + _modes + neuron_config = NeuronConfig( + modes=_modes, + profile_output_dir=rank_dir, + max_events_per_nc=10_000_000, + capture_enabled_for_nc="0,1", + ) + exporter = NeuronProfiler(neuron_config) + + if rank == 0: + _tag = "RUNTIME-only" if _runtime_only else "DEVICE+RUNTIME" + log( + rank, + f"=== PROFILING ({_tag}) active={args.profile_active} dir={profile_dir} ===", + ) + dist.barrier() + # wait=1, warmup=1 to settle, then capture `active` steady-state steps. + sched = schedule(wait=1, warmup=1, active=args.profile_active, repeat=1) + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1], + experimental_config=neuron_config, + schedule=sched, + on_trace_ready=exporter.export_trace, + record_shapes=True, + ) as prof: + for _ in range(2 + args.profile_active): + one_step(collect=False) + torch.neuron.synchronize() + prof.step() + # Also drop a per-rank framework chrome trace next to the session. + try: + prof.export_chrome_trace(os.path.join(rank_dir, f"framework_rank{rank}.json")) + except Exception as e: # noqa: BLE001 + log(rank, f"chrome-trace export skipped: {e}") + dist.barrier() + if rank == 0: + log(rank, f"=== PROFILE session dir: {neuron_config.get_session_dir()} ===") + log(rank, "=== PROFILE DONE ===") + + +if __name__ == "__main__": + main() diff --git a/torchtitan_neuron/tests/nki_blockwise_moe_cte.py b/torchtitan_neuron/tests/nki_blockwise_moe_cte.py new file mode 100644 index 00000000000..31cf6379a9d --- /dev/null +++ b/torchtitan_neuron/tests/nki_blockwise_moe_cte.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Shared handle for the pure blockwise MoE-CTE NKI kernel (no router). + +``moe_cte`` (nkilib.core.moe.moe_cte) is the standalone blockwise expert matmul: it takes +block metadata (token_position_to_id, block_to_expert, block_size) as inputs and tiles +internally by block_size, so it compiles where a single large ``aten::_grouped_mm`` over +the same rows overflows the neuronx-cc 5M-instruction limit (NCC_EBVF030). + +Exposes one wrapped kernel and a default spec so the standalone blockwise bench +and the neuron_ep dispatcher bench share one expert-compute definition. +""" + +from __future__ import annotations + +from nkilib.core.moe.moe_cte import ( + moe_cte, + MoECTEImplementation, + MoECTESpec, + ShardOnIConfig, +) + +from torch_neuronx.nki_hop import wrap_nki + +# wrap_nki handle; index with [(lnc,)] at the callsite. +wrapped_moe_cte = wrap_nki(moe_cte) + + +def default_spec(): + """Simplest working spec: I-sharding, defaults. bf16 TRN2/TRN3 path.""" + return MoECTESpec( + implementation=MoECTEImplementation.shard_on_i, + shard_on_I=ShardOnIConfig(), + ) diff --git a/torchtitan_neuron/tests/test_attention_masks.py b/torchtitan_neuron/tests/test_attention_masks.py new file mode 100644 index 00000000000..c57643abae5 --- /dev/null +++ b/torchtitan_neuron/tests/test_attention_masks.py @@ -0,0 +1,98 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Unit tests for GPT-OSS attention mask construction (get_attention_masks). + +GPT-OSS alternates full-attention and sliding-window-attention layers. Its +``GptOssModel.get_attention_masks`` builds a dict of block masks on the flex path: +always a ``basic_mask``, plus a ``sliding_window_mask`` when any layer requests a +window. This covers that branching logic and the composed mask_mods. + +``create_attention_mask`` (which compiles ``create_block_mask``) is patched out so +the tests exercise GPT-OSS's own logic -- the mask_mod composition and per-mask +selection -- without depending on block-mask compilation (device/backend-specific). +The composed mask_mod is captured and evaluated directly to verify the window +restriction. ``get_attention_masks`` reads only ``self.config``, so we drive it with +the debug config rather than instantiating the full model. +""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +import pytest + +import torch + +# Importing the gpt_oss model pulls torchtitan.distributed, which needs the Neuron +# stack. Gate so collection works on a host without it. +pytest.importorskip("spmd_types") + +from torchtitan.models.gpt_oss import model_registry # noqa: E402 +from torchtitan.models.gpt_oss.model import GptOssModel # noqa: E402 + + +def _flex_debug_config(): + # attn_backend="flex" so inner_attention is FlexAttention.Config (the debug + # model puts a sliding window on even-indexed layers). + return model_registry("debugmodel", attn_backend="flex").model + + +class TestGptOssAttentionMasks(unittest.TestCase): + S = 256 + + def _build_masks(self, cfg): + """Run get_attention_masks with create_attention_mask stubbed to record the + composed mask_mod per named mask (returns the mask_mod itself as the mask).""" + positions = torch.arange(self.S).unsqueeze(0) # [B=1, S] + + def _capture(mask_mod, B, H, q_len, kv_len, **kwargs): + return mask_mod # stand in for the BlockMask; keep the mask_mod itself + + with mock.patch( + "torchtitan.models.gpt_oss.model.create_attention_mask", _capture + ): + return GptOssModel.get_attention_masks( + SimpleNamespace(config=cfg), positions + ) + + def test_flex_builds_basic_and_sliding_window_masks(self): + masks = self._build_masks(_flex_debug_config()) + self.assertIsInstance(masks, dict) + self.assertEqual(set(masks), {"basic_mask", "sliding_window_mask"}) + + def test_only_basic_mask_when_no_layer_uses_sliding_window(self): + cfg = _flex_debug_config() + for layer in cfg.layers: + layer.attention.sliding_window_size = None + masks = self._build_masks(cfg) + self.assertEqual(set(masks), {"basic_mask"}) + + def test_sliding_window_mask_restricts_range(self): + cfg = _flex_debug_config() + # The debug model's windowed layers use sliding_window_size=128. + window = next( + layer.attention.sliding_window_size + for layer in cfg.layers + if layer.attention.sliding_window_size is not None + ) + composed_mask_mod = self._build_masks(cfg)["sliding_window_mask"] + + def allowed(q_idx: int, kv_idx: int) -> bool: + zero = torch.tensor(0) + return bool( + composed_mask_mod(zero, zero, torch.tensor(q_idx), torch.tensor(kv_idx)) + ) + + # Causal + within-window (q - kv < window) -> attention allowed. + self.assertTrue(allowed(200, 200 - (window - 1))) + # Causal but beyond the window (q - kv >= window) -> masked out. + self.assertFalse(allowed(200, 200 - window)) + # Future token (kv > q) -> masked by causality. + self.assertFalse(allowed(200, 201)) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan_neuron/tests/test_capability.py b/torchtitan_neuron/tests/test_capability.py new file mode 100644 index 00000000000..c85088a71d4 --- /dev/null +++ b/torchtitan_neuron/tests/test_capability.py @@ -0,0 +1,80 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Unit tests for torchtitan.tools.utils.has_neuron_capability. + +The off-device and simulated-trn3 cases run anywhere (they do not need real +Neuron hardware or the neuron custom ops). The real-hardware case is gated on +trn3 because MXFP8 is a trn3-only capability. +""" + +import sys +import types +import unittest +from unittest import mock + +from torchtitan.tools.utils import has_neuron_capability + + +class TestHasNeuronCapability(unittest.TestCase): + def test_returns_false_without_torch_neuronx(self): + # When torch_neuronx can't be imported, mxfp8 is unsupported -- forced + # here (sys.modules entry -> None makes the import raise) so it holds on + # any platform, including a real trn3 pod where it would otherwise be True. + with mock.patch.dict(sys.modules, {"torch_neuronx": None}): + self.assertFalse(has_neuron_capability("mxfp8")) + + def test_unknown_feature_returns_false(self): + # Only "mxfp8" is a real capability today; anything else is False, + # including on hardware. + self.assertFalse(has_neuron_capability("fp4")) + self.assertFalse(has_neuron_capability("nvfp4")) + + def test_feature_name_is_case_insensitive(self): + # "MXFP8" and "mxfp8" resolve identically (the .lower() path) on any + # platform -- both False off-device, both True on trn3. + self.assertEqual(has_neuron_capability("MXFP8"), has_neuron_capability("mxfp8")) + + def _fake_torch_neuronx(self, platform_target: str) -> types.ModuleType: + mod = types.ModuleType("torch_neuronx") + utils_mod = types.ModuleType("torch_neuronx.utils") + utils_mod.get_platform_target = lambda: platform_target + mod.utils = utils_mod + return mod + + def test_simulated_trn3_returns_true(self): + fake = self._fake_torch_neuronx("trn3") + with mock.patch.dict( + sys.modules, + {"torch_neuronx": fake, "torch_neuronx.utils": fake.utils}, + ): + self.assertTrue(has_neuron_capability("mxfp8")) + + def test_simulated_trn3u_returns_true(self): + # startswith("trn3") also matches later trn3 variants. + fake = self._fake_torch_neuronx("trn3u") + with mock.patch.dict( + sys.modules, + {"torch_neuronx": fake, "torch_neuronx.utils": fake.utils}, + ): + self.assertTrue(has_neuron_capability("mxfp8")) + + def test_simulated_trn2_returns_false(self): + fake = self._fake_torch_neuronx("trn2") + with mock.patch.dict( + sys.modules, + {"torch_neuronx": fake, "torch_neuronx.utils": fake.utils}, + ): + self.assertFalse(has_neuron_capability("mxfp8")) + + @unittest.skipUnless( + has_neuron_capability("mxfp8"), "requires a real Neuron trn3 device" + ) + def test_real_trn3_returns_true(self): + self.assertTrue(has_neuron_capability("mxfp8")) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan_neuron/tests/test_compile_backend.py b/torchtitan_neuron/tests/test_compile_backend.py new file mode 100644 index 00000000000..0878d2793ab --- /dev/null +++ b/torchtitan_neuron/tests/test_compile_backend.py @@ -0,0 +1,47 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +from typing import Any, cast +from unittest import mock + +import torch + +from torchtitan_neuron.compile_backend import install_neuron_compile_backend + + +def test_install_neuron_compile_backend_routes_only_backendless_calls(monkeypatch): + original_compile = mock.Mock(return_value=object()) + original_compile._neuron_wrapped = False + monkeypatch.setattr(torch, "compile", original_compile) + + install_neuron_compile_backend() + + def model(value): + return value + + options = {"example": True} + wrapped_compile = cast(Any, torch.compile) + wrapped_compile(model, options=options) + original_compile.assert_called_once_with( + model, + options=options, + backend="neuron", + ) + + original_compile.reset_mock() + wrapped_compile(model, backend="custom") + original_compile.assert_called_once_with(model, backend="custom") + + +def test_install_neuron_compile_backend_is_idempotent(monkeypatch): + original_compile = mock.Mock(return_value=object()) + original_compile._neuron_wrapped = False + monkeypatch.setattr(torch, "compile", original_compile) + + install_neuron_compile_backend() + wrapped_compile = torch.compile + install_neuron_compile_backend() + + assert torch.compile is wrapped_compile diff --git a/torchtitan_neuron/tests/test_dense_config_registry.py b/torchtitan_neuron/tests/test_dense_config_registry.py new file mode 100644 index 00000000000..78cd1d1c0cc --- /dev/null +++ b/torchtitan_neuron/tests/test_dense_config_registry.py @@ -0,0 +1,143 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +import pytest + + +pytest.importorskip("spmd_types") + +from torchtitan.components.loss import ( # noqa: E402 + ChunkedLossWrapper, + CrossEntropyLoss, +) +from torchtitan.distributed.activation_checkpoint import FullAC # noqa: E402 +from torchtitan.experiments.graph_trainer.chunked_loss import ( # noqa: E402 + ChunkedLossWrapperWithParamGrads, +) +from torchtitan.experiments.graph_trainer.configs import ( # noqa: E402 + GraphTrainerCompileConfig, +) +from torchtitan.experiments.graph_trainer.qwen3.model import ( # noqa: E402 + GraphTrainerQwen3Model, +) +from torchtitan.models.common.attention import ScaledDotProductAttention # noqa: E402 +from torchtitan.models.qwen3 import ( # noqa: E402 + config_registry as core_qwen3_config_registry, +) + +from torchtitan_neuron.experiments.graph_trainer.qwen3 import ( # noqa: E402 + config_registry as graph_trainer_qwen3_config_registry, +) +from torchtitan_neuron.models.qwen3 import ( # noqa: E402 + config_registry as qwen3_config_registry, +) + + +def _lr(config): + return config.optimizer.param_groups[0].optimizer_kwargs["lr"] + + +def _assert_chunked_cross_entropy_loss(config): + assert isinstance(config.loss, ChunkedLossWrapper.Config) + assert isinstance(config.loss.loss_fn, CrossEntropyLoss.Config) + + +def _assert_sdpa(config): + assert isinstance( + config.model_spec.model.layers[0].attention.inner_attention, + ScaledDotProductAttention.Config, + ) + + +def test_neuron_dense_recipes_are_not_registered_in_core(): + assert not hasattr( + core_qwen3_config_registry, + "qwen3_32b_tp4_fsdp16_fullac_compile", + ) + + +def test_qwen3_8b_dense_eager_contract(): + config = qwen3_config_registry.qwen3_8b_tp1_fsdp64_fullac_eager() + + assert config.model_spec.name == "qwen3" + assert config.model_spec.flavor == "8B" + _assert_chunked_cross_entropy_loss(config) + _assert_sdpa(config) + assert config.parallelism.tensor_parallel_degree == 1 + assert config.parallelism.data_parallel_shard_degree == -1 + assert config.training.local_batch_size == 1 + assert config.training.global_batch_size == 64 + assert _lr(config) == 8e-4 + assert isinstance(config.activation_checkpoint, FullAC.Config) + assert config.compile.enable is False + assert config.override.imports == [] + + +def test_qwen3_8b_dense_compile_contract(): + config = qwen3_config_registry.qwen3_8b_tp1_fsdp64_fullac_compile() + + assert config.model_spec.name == "qwen3" + assert config.model_spec.flavor == "8B" + _assert_chunked_cross_entropy_loss(config) + _assert_sdpa(config) + assert config.parallelism.tensor_parallel_degree == 1 + assert config.parallelism.data_parallel_shard_degree == -1 + assert config.training.local_batch_size == 1 + assert config.training.global_batch_size == 64 + assert _lr(config) == 8e-4 + assert isinstance(config.activation_checkpoint, FullAC.Config) + assert config.compile.enable is True + assert config.compile.backend == "neuron" + assert config.compile.components == ["model", "loss", "grad_clip"] + assert config.override.imports == [] + + +def test_simple_fsdp_qwen3_8b_config_contract(): + config = ( + graph_trainer_qwen3_config_registry.simple_fsdp_qwen3_8b_tp1_fsdp64_fullac() + ) + + assert config.model_spec.name == "simple_fsdp/qwen3" + assert config.model_spec.flavor == "8B" + assert isinstance(config.model_spec.model, GraphTrainerQwen3Model.Config) + _assert_sdpa(config) + + assert config.parallelism.tensor_parallel_degree == 1 + assert config.parallelism.data_parallel_shard_degree == -1 + assert config.training.local_batch_size == 1 + assert config.training.global_batch_size == 64 + assert isinstance(config.activation_checkpoint, FullAC.Config) + assert isinstance(config.compile, GraphTrainerCompileConfig) + assert config.compile.enable is True + assert config.compile.backend == "neuron" + assert config.compile.mode == "jit" + assert config.compile.components == ["model", "loss", "grad_clip"] + + assert isinstance(config.loss, ChunkedLossWrapperWithParamGrads.Config) + assert isinstance(config.loss.loss_fn, CrossEntropyLoss.Config) + assert config.override.imports == [] + assert all( + layer.attention.rope.use_prefix_cache_for_positions + for layer in config.model_spec.model.layers + ) + + +def test_qwen3_32b_dense_compile_contract(): + config = qwen3_config_registry.qwen3_32b_tp4_fsdp16_fullac_compile() + + assert config.model_spec.name == "qwen3" + assert config.model_spec.flavor == "32B" + _assert_chunked_cross_entropy_loss(config) + _assert_sdpa(config) + assert config.parallelism.tensor_parallel_degree == 4 + assert config.parallelism.data_parallel_shard_degree == -1 + assert config.training.local_batch_size == 1 + assert config.training.global_batch_size == 16 + assert _lr(config) == 8e-4 + assert isinstance(config.activation_checkpoint, FullAC.Config) + assert config.compile.enable is True + assert config.compile.backend == "neuron" + assert config.compile.components == ["model", "loss", "grad_clip"] + assert config.override.imports == [] diff --git a/torchtitan_neuron/tests/test_ep_dispatcher.py b/torchtitan_neuron/tests/test_ep_dispatcher.py new file mode 100644 index 00000000000..deb83c1e73e --- /dev/null +++ b/torchtitan_neuron/tests/test_ep_dispatcher.py @@ -0,0 +1,929 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Unit tests for the neuron_ep token dispatcher layout logic. + +These cover the CPU-computable, non-kernel parts of NeuronEPTokenDispatcher: +packed routing-vector forwarding, capacity validation, mesh-wiring validation, data +provenance across ranks, the EP=64 metadata contract, and the EP=1 local-dispatch +fallback. The NKI ``permute_a2av`` / ``unpermute_a2av`` kernels themselves are +exercised on Trainium hardware; they are not invoked here. +""" + +import unittest +import unittest.mock + +import torch + +from torchtitan.distributed.neuron_ep import gmm_align +from torchtitan.models.common.token_dispatcher import ( + NeuronEPDispatchMetadata, + NeuronEPTokenDispatcher, +) + + +def _make_dispatcher(num_experts: int, top_k: int, ep_size: int | None): + cfg = NeuronEPTokenDispatcher.Config(num_experts=num_experts, top_k=top_k) + dispatcher = NeuronEPTokenDispatcher(cfg) + if ep_size is None: + dispatcher.ep_mesh = None + else: + mock_mesh = unittest.mock.MagicMock() + mock_mesh.size.return_value = ep_size + dispatcher.ep_mesh = mock_mesh + dispatcher._ep_ranks = [list(range(ep_size))] + return dispatcher + + +def _counts_from_ids(topk_ids_TK: torch.Tensor, num_experts: int) -> torch.Tensor: + counts = torch.zeros(num_experts, dtype=torch.int64) + for e in topk_ids_TK.reshape(-1): + counts[int(e)] += 1 + return counts + + +def setUpModule(): # noqa: N802 + # The dispatcher tests use an independent CPU oracle at the NKI boundary. + # Tile-prefix construction itself is covered by on-device kernel tests. + def packed_indices(ids, counts): + expected_counts = torch.bincount(ids.long(), minlength=counts.numel()) + torch.testing.assert_close(counts.flatten().long(), expected_counts) + return torch.argsort(ids, stable=True).to(torch.int32) + + patcher = unittest.mock.patch( + "torchtitan.models.common.token_dispatcher.build_packed_source_indices", + side_effect=packed_indices, + ) + patcher.start() + unittest.addModuleCleanup(patcher.stop) + + +class TestWireMeshesValidation(unittest.TestCase): + def test_accepts_sp(self): + """TP/SP is supported: wire_meshes records the SP coordinates.""" + disp = _make_dispatcher(128, 8, ep_size=None) + tp_mesh = unittest.mock.MagicMock() + tp_mesh.size.return_value = 4 + tp_mesh._sym_get_coordinate.return_value = 2 + disp.wire_meshes(ep_mesh=None, tp_mesh=tp_mesh) + self.assertEqual(disp.sp_size, 4) + self.assertEqual(disp.sp_rank, 2) + + def test_rejects_indivisible_experts(self): + disp = _make_dispatcher(num_experts=127, top_k=8, ep_size=None) + ep_mesh = unittest.mock.MagicMock() + ep_mesh.size.return_value = 4 + with self.assertRaisesRegex(ValueError, "divisible"): + disp.wire_meshes(ep_mesh=ep_mesh, tp_mesh=None) + + +class TestEP1Fallback(unittest.TestCase): + """With ep_mesh None, dispatch/combine reduce to LocalTokenDispatcher.""" + + def test_dispatch_combine_roundtrip(self): + torch.manual_seed(0) + num_experts, top_k = 8, 2 + T, D = 6, 16 + disp = _make_dispatcher(num_experts, top_k, ep_size=None) + + x_TD = torch.randn(T, D) + scores_TK = torch.rand(T, top_k) + ids_TK = torch.randint(0, num_experts, (T, top_k)) + counts_E = _counts_from_ids(ids_TK, num_experts) + + routed_RD, counts_e, md = disp.dispatch(x_TD, scores_TK, ids_TK, counts_E) + # No all-to-all: R == N == T*top_k. + self.assertEqual(routed_RD.shape[0], T * top_k) + + # Identity experts -> combine reproduces the score-weighted top-k sum. + out_TD = disp.combine( + routed_RD, + md, + x_TD, + num_local_tokens_after_padding=0, + local_seq_len_after_padding=0, + ) + self.assertEqual(out_TD.shape, x_TD.shape) + expected = (x_TD.unsqueeze(1) * scores_TK.unsqueeze(-1)).sum(dim=1) + torch.testing.assert_close(out_TD, expected, rtol=1e-4, atol=1e-4) + + +def _make_compact_dispatcher( + num_experts, + top_k, + ep_size, + *, + capacity_per_expert=None, +): + """Dispatcher configured for the supported packed path.""" + cfg = NeuronEPTokenDispatcher.Config( + num_experts=num_experts, + top_k=top_k, + capacity_per_expert=capacity_per_expert, + ) + disp = NeuronEPTokenDispatcher(cfg) + mock_mesh = unittest.mock.MagicMock() + mock_mesh.size.return_value = ep_size + disp.ep_mesh = mock_mesh + disp._ep_ranks = [list(range(ep_size))] + return disp + + +class TestOptimizedCompactTransmission(unittest.TestCase): + """Compact (variable-count) dispatch for num_local_experts == 1. + + Sends only real per-destination rows and receives them at cumulative source + offsets in an FSDP/global-token-bound capacity buffer. + + Requires nle == 1 (EP == num_experts): only then are the kernel's per-source + recv_counts also per-(source, expert), which the expert buffer needs. + """ + + def test_packed_capacity_matches_fsdp_batch_sequence_bound(self): + batch_size = 2 + local_seq_len = 512 + sp_size = 4 + fsdp_size = 2 + ep_size = fsdp_size * sp_size + disp = _make_compact_dispatcher(ep_size, 4, ep_size) + capacity = disp._packed_capacity(batch_size * local_seq_len, ep_size, 1) + expected = fsdp_size * batch_size * (local_seq_len * sp_size) + self.assertEqual(capacity, expected) + + def test_packed_capacity_rounds_model_bound_for_gmm(self): + ep_size, top_k, num_local_tokens = 8, 4, 129 + disp = _make_compact_dispatcher(ep_size, top_k, ep_size) + capacity = disp._packed_capacity(num_local_tokens, ep_size, 1) + unaligned_bound = num_local_tokens * ep_size + self.assertGreaterEqual(capacity, unaligned_bound) + self.assertLess(capacity - unaligned_bound, gmm_align()) + self.assertEqual(capacity % gmm_align(), 0) + + def test_packed_capacity_is_independent_of_top_k(self): + num_local_tokens = 1024 + ep_size = 64 + capacities = [ + _make_compact_dispatcher(ep_size, top_k, ep_size)._packed_capacity( + num_local_tokens, ep_size, 1 + ) + for top_k in (1, 2, 4, 8) + ] + self.assertEqual(capacities, [num_local_tokens * ep_size] * 4) + + def test_capacity_override_retains_per_source_semantics(self): + ep_size = 8 + disp = _make_compact_dispatcher( + ep_size, + 4, + ep_size, + capacity_per_expert=1024, + ) + self.assertEqual(disp._packed_capacity(1024, ep_size, 1), ep_size * 1024) + + def test_capacity_override_rejects_truncating_bound(self): + ep_size = 8 + disp = _make_compact_dispatcher( + ep_size, + 4, + ep_size, + capacity_per_expert=127, + ) + with self.assertRaisesRegex(ValueError, "model capacity"): + disp._packed_capacity(1024, ep_size, 1) + + def test_dispatch_rejects_mismatched_routing_shapes(self): + disp = _make_compact_dispatcher(4, 2, 4) + x = torch.randn(3, 8) + ids = torch.tensor([[0, 1], [1, 2], [2, 3]]) + scores = torch.randn(3, 1) + with self.assertRaisesRegex(ValueError, "matching top-k"): + disp.dispatch(x, scores, ids, torch.zeros(4, dtype=torch.int32)) + + def test_dispatch_rejects_routing_width_that_differs_from_config(self): + disp = _make_compact_dispatcher(4, 2, 4) + x = torch.randn(3, 8) + ids = torch.tensor([[0, 1, 2], [1, 2, 3], [0, 2, 3]]) + scores = torch.randn(3, 3) + with self.assertRaisesRegex(ValueError, "routing shape \\(3, 2\\)"): + disp.dispatch(x, scores, ids, torch.zeros(4, dtype=torch.int32)) + + def test_dispatch_rejects_routing_token_count_that_differs_from_input(self): + disp = _make_compact_dispatcher(4, 2, 4) + x = torch.randn(3, 8) + ids = torch.tensor([[0, 1], [1, 2]]) + scores = torch.randn(2, 2) + with self.assertRaisesRegex(ValueError, "routing shape \\(3, 2\\)"): + disp.dispatch(x, scores, ids, torch.zeros(4, dtype=torch.int32)) + + def test_combine_rejects_output_that_does_not_match_compact_capacity(self): + ep_size, capacity, hidden = 4, 3, 8 + disp = _make_compact_dispatcher( + num_experts=ep_size, + top_k=2, + ep_size=ep_size, + ) + metadata = NeuronEPDispatchMetadata( + token_indices_experts_sorted_N=torch.empty(0, dtype=torch.int64), + topk_scores_experts_sorted_N=torch.empty(0), + packed_source_indices_N=torch.empty(0, dtype=torch.int32), + send_counts_1E=torch.zeros(1, ep_size, dtype=torch.int32), + recv_counts_1E=torch.zeros(1, ep_size, dtype=torch.int32), + compact_capacity=capacity, + ) + routed_output = torch.empty(capacity - 1, hidden) + + with self.assertRaisesRegex(ValueError, "capacity"): + disp.combine( + routed_output, + metadata, + torch.empty(1, hidden), + num_local_tokens_after_padding=1, + local_seq_len_after_padding=1, + ) + + def _packed_permute_model( + self, routed_ND, send_vector, send_cnt, ep_size, n, capacity, D + ): + """CPU model of explicit-displacement packed dispatch.""" + # send_rows[(s, d)] = the real rows source s sends to dest d. + send_rows = {} + for s in range(ep_size): + for d in range(ep_size): + cnt = int(send_cnt[s][0, d]) + rows = [] + offset = int(send_cnt[s][0, :d].sum()) + for slot in range(cnt): + src_row = int(send_vector[s][offset + slot]) + self.assertLess(src_row, n) + rows.append(routed_ND[s][src_row]) + send_rows[(s, d)] = rows + recv = {} + for d in range(ep_size): + recv_counts = [len(send_rows[(s, d)]) for s in range(ep_size)] + total = sum(recv_counts) + self.assertLessEqual(total, capacity) + buf = torch.zeros(capacity, D) + offset = 0 + for s in range(ep_size): + for j, row in enumerate(send_rows[(s, d)]): + buf[offset + j] = row + offset += recv_counts[s] + recv[d] = ( + buf, + torch.tensor(recv_counts, dtype=torch.int32).view(1, ep_size), + ) + return recv + + def test_packed_routing_and_receive_offsets(self): + torch.manual_seed(0) + # nle == 1: num_experts == ep_size. + num_experts = ep_size = 8 + top_k = 4 + T, D = 12, 16 + n = T * top_k + capacity = 2 * n + + disp = [ + _make_compact_dispatcher(num_experts, top_k, ep_size) + for _ in range(ep_size) + ] + x = [torch.randn(T, D) for _ in range(ep_size)] + # Distinct top_k experts per token (precondition) so each expert gets <= T rows. + ids = [ + torch.stack([torch.randperm(num_experts)[:top_k] for _ in range(T)]) + for _ in range(ep_size) + ] + routed_ND = [x[r].repeat_interleave(top_k, dim=0) for r in range(ep_size)] + + # Peer counts partition the vector, preserving each destination's arrival order. + send_vector, send_cnt = [], [] + for r in range(ep_size): + sc, si = disp[r]._build_packed_routing( + ids[r], _counts_from_ids(ids[r], ep_size) + ) + send_vector.append(si) + send_cnt.append(sc) + self.assertEqual(int(sc.sum()), n) + flat_ids = ids[r].reshape(-1) + for d in range(ep_size): + cnt = int(sc[0, d]) + self.assertEqual(cnt, int((flat_ids == d).sum())) # dest == expert + offset = int(sc[0, :d].sum()) + expected = torch.where(flat_ids == d)[0].to(torch.int32) + torch.testing.assert_close(si[offset : offset + cnt], expected) + torch.testing.assert_close( + si.sort().values, torch.arange(n, dtype=torch.int32) + ) + + # 2. Incoming source chunks are adjacent at cumulative receive offsets. + recv = self._packed_permute_model( + routed_ND, send_vector, send_cnt, ep_size, n, capacity, D + ) + for d in range(ep_size): + recv_buf, recv_counts_1E = recv[d] + rc = recv_counts_1E.reshape(-1) + offset = 0 + for s in range(ep_size): + cs = int(rc[s]) + for j in range(cs): + source_offset = int(send_cnt[s][0, :d].sum()) + expected = routed_ND[s][int(send_vector[s][source_offset + j])] + torch.testing.assert_close(recv_buf[offset + j], expected) + offset += cs + self.assertTrue(bool((recv_buf[offset:] == 0).all())) + + def test_compact_end_to_end_matches_reference(self): + # Full dispatch->expert->combine round-trip under the packed receive model, + # with distinct per-expert weights so a misroute changes the value. + torch.manual_seed(3) + num_experts = ep_size = 8 + top_k = 4 + T, D = 12, 16 + n = T * top_k + capacity = 2 * n + E = num_experts + + def expert_fn(row, g): + return row * (1.0 + g / E) + + disp = [ + _make_compact_dispatcher(num_experts, top_k, ep_size) + for _ in range(ep_size) + ] + x = [torch.randn(T, D) for _ in range(ep_size)] + ids = [ + torch.stack([torch.randperm(num_experts)[:top_k] for _ in range(T)]) + for _ in range(ep_size) + ] + scores = [torch.rand(T, top_k) for _ in range(ep_size)] + routed_ND = [x[r].repeat_interleave(top_k, dim=0) for r in range(ep_size)] + tok_idx = [torch.arange(n) // top_k for _ in range(ep_size)] + tok_scores = [scores[r].reshape(-1) for r in range(ep_size)] + + send_vector, send_cnt = [], [] + for r in range(ep_size): + sc, si = disp[r]._build_packed_routing( + ids[r], _counts_from_ids(ids[r], ep_size) + ) + send_vector.append(si) + send_cnt.append(sc) + + recv = self._packed_permute_model( + routed_ND, send_vector, send_cnt, ep_size, n, capacity, D + ) + + expert_out_packed = {} + for d in range(ep_size): + recv_buf, recv_counts_1E = recv[d] + expert_out_packed[d] = (expert_fn(recv_buf, d), recv_counts_1E) + + # Reverse A2AV reads this origin's chunk at its cumulative source offset. + out = [torch.zeros(T, D) for _ in range(ep_size)] + for o in range(ep_size): + combined_ND = torch.zeros(n, D) + for d in range(ep_size): + out_packed, recv_counts_1E = expert_out_packed[d] + cnt = int(send_cnt[o][0, d]) + offset = int(recv_counts_1E[0, :o].sum()) + for slot in range(cnt): + source_offset = int(send_cnt[o][0, :d].sum()) + src_row = int(send_vector[o][source_offset + slot]) + combined_ND[src_row] = out_packed[offset + slot] + combined_ND = combined_ND * tok_scores[o].reshape(-1, 1) + out[o] = torch.zeros(T, D).index_add(0, tok_idx[o], combined_ND) + + for o in range(ep_size): + expected = torch.zeros(T, D) + for t in range(T): + for kk in range(top_k): + g = int(ids[o][t, kk]) + expected[t] += scores[o][t, kk] * expert_fn(x[o][t], g) + torch.testing.assert_close(out[o], expected, rtol=1e-4, atol=1e-4) + + def test_compact_rejects_nle_gt_1(self): + disp = _make_compact_dispatcher(num_experts=128, top_k=8, ep_size=64) # nle=2 + T = 16 + x = torch.randn(T, 8) + ids = torch.stack([torch.randperm(128)[:8] for _ in range(T)]) + scores = torch.rand(T, 8) + counts = _counts_from_ids(ids, 128) + with self.assertRaisesRegex(ValueError, "num_local_experts == 1"): + disp.dispatch(x, scores, ids, counts) + + def test_ep_below_kernel_minimum_is_rejected(self): + ep_size = 4 + disp = _make_compact_dispatcher(ep_size, 2, ep_size) + T = 4 + x = torch.randn(T, 8) + ids = torch.stack([torch.randperm(ep_size)[:2] for _ in range(T)]) + scores = torch.rand(T, 2) + counts = _counts_from_ids(ids, ep_size) + with self.assertRaisesRegex(ValueError, "ep_size >= 8"): + disp.dispatch(x, scores, ids, counts) + + def test_ep_without_complete_four_rank_devices_is_rejected(self): + ep_size = 10 + disp = _make_compact_dispatcher(ep_size, 2, ep_size) + T = 4 + x = torch.randn(T, 8) + ids = torch.stack([torch.randperm(ep_size)[:2] for _ in range(T)]) + scores = torch.rand(T, 2) + counts = _counts_from_ids(ids, ep_size) + with self.assertRaisesRegex(ValueError, "complete four-rank"): + disp.dispatch(x, scores, ids, counts) + + +class TestSequenceParallelCombine(unittest.TestCase): + """combine() under TP/SP (sp_size > 1) scatters to GLOBAL token positions. + + With SP each rank holds a 1/sp_size slice of the sequence, so combine must write + into a global-length buffer at the offsets given by ``_sp_global_token_indices`` + (blocked layout: local position p of batch row b -> b*global_seq + sp_rank*L + p), + leaving every other rank's slice zero. The downstream Partial -> Shard(seq) + redistribution reassembles the slices, so a rank writing outside its own slice (or + a local-length buffer) would silently corrupt the sequence. + + NOTE on coverage: the ``_run`` helper below drives the EP=1 (``ep_mesh=None``) path, + which returns from ``combine`` via ``LocalTokenDispatcher`` BEFORE the sp_size > 1 + branch -- so it pins the delegation contract only, NOT the SP scatter. + ``TestSequenceParallelCombineEPPath`` covers the real branch. + """ + + def _run(self, *, sp_size, sp_rank, B, local_L, D=8, num_experts=8, top_k=2): + torch.manual_seed(0) + disp = _make_dispatcher(num_experts, top_k, ep_size=None) + tp_mesh = unittest.mock.MagicMock() + tp_mesh.size.return_value = sp_size + tp_mesh._sym_get_coordinate.return_value = sp_rank + disp.wire_meshes(ep_mesh=None, tp_mesh=tp_mesh) + self.assertEqual(disp.sp_size, sp_size) + + t_local = B * local_L + x_TD = torch.randn(t_local, D) + scores_TK = torch.rand(t_local, top_k) + ids_TK = torch.randint(0, num_experts, (t_local, top_k)) + counts_E = _counts_from_ids(ids_TK, num_experts) + + routed_RD, _counts_e, md = disp.dispatch(x_TD, scores_TK, ids_TK, counts_E) + out = disp.combine( + routed_RD, + md, + x_TD, + num_local_tokens_after_padding=t_local, + local_seq_len_after_padding=local_L, + ) + expected_local = (x_TD.unsqueeze(1) * scores_TK.unsqueeze(-1)).sum(dim=1) + return out, expected_local + + def test_ep1_sp_falls_back_to_local_shape(self): + """EP=1 delegates to LocalTokenDispatcher: local-length, no SP offset. + + Guards the delegation contract -- the SP-aware global scatter lives in the EP + path; the EP=1 fallback stays local-length (its own scatter target is x_TD). + """ + out, expected = self._run(sp_size=4, sp_rank=1, B=2, local_L=3) + self.assertEqual(out.shape, expected.shape) + torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4) + + def test_global_scatter_positions_blocked_layout(self): + """The SP remap places local ids at the expected global rows, disjointly.""" + sp_size, sp_rank, B, local_L = 4, 2, 2, 3 + disp = _make_dispatcher(8, 2, ep_size=None) + tp_mesh = unittest.mock.MagicMock() + tp_mesh.size.return_value = sp_size + tp_mesh._sym_get_coordinate.return_value = sp_rank + disp.wire_meshes(ep_mesh=None, tp_mesh=tp_mesh) + + local_ids = torch.arange(B * local_L) + global_ids = disp._sp_global_token_indices(local_ids, local_L) + global_seq = local_L * sp_size + expected = torch.tensor( + [ + (i // local_L) * global_seq + sp_rank * local_L + (i % local_L) + for i in range(B * local_L) + ] + ) + torch.testing.assert_close(global_ids, expected) + # Every index lands inside this rank's own slice of each batch row, and within + # the global buffer -- i.e. ranks write disjoint regions. + self.assertTrue(int(global_ids.max()) < B * global_seq) + for i, g in enumerate(global_ids.tolist()): + pos_in_row = g % global_seq + self.assertGreaterEqual(pos_in_row, sp_rank * local_L) + self.assertLess(pos_in_row, (sp_rank + 1) * local_L) + self.assertEqual(g // global_seq, i // local_L) # batch row preserved + + def test_sp_rank0_matches_no_sp_prefix(self): + """At sp_rank=0 the global ids of batch row 0 are the identity prefix.""" + disp = _make_dispatcher(8, 2, ep_size=None) + tp_mesh = unittest.mock.MagicMock() + tp_mesh.size.return_value = 2 + tp_mesh._sym_get_coordinate.return_value = 0 + disp.wire_meshes(ep_mesh=None, tp_mesh=tp_mesh) + ids = torch.arange(4) # B=1, local_L=4 + torch.testing.assert_close(disp._sp_global_token_indices(ids, 4), ids) + + +class TestSequenceParallelCombineEPPath(unittest.TestCase): + """combine() SP branch on the REAL EP path (ep_mesh set), with kernels stubbed. + + This is the test that actually executes the sp_size > 1 code: the global-length + output buffer and the ``_sp_global_token_indices`` remap of the final scatter. The + EP=1 tests above return early via LocalTokenDispatcher and never reach it, so a + regression to ``torch.zeros_like(x_TD)`` + raw local indices would not be caught + without this. + + The packed NKI dispatch and combine kernels are replaced with CPU stubs: + dispatch packs this rank's rows by destination, and combine restores their + original order. That leaves the index/scatter logic under test without + needing Trainium or ``torch_neuronx``. ``ep_mesh`` / ``_ep_ranks`` are set + directly rather than through ``wire_meshes``, which would call + ``build_ep_replica_group_ranks`` (needs + ``torch_neuronx.distributed.mesh_registry``). + """ + + def _make(self, *, num_experts, top_k, ep_size, sp_size, sp_rank): + disp = _make_compact_dispatcher(num_experts, top_k, ep_size) + disp.sp_size = sp_size + disp.sp_rank = sp_rank + return disp + + def _run_ep_sp(self, *, sp_size, sp_rank, B, local_L, D=8, num_experts=8, top_k=2): + """Drive dispatch+combine through the EP path with stubbed kernels.""" + torch.manual_seed(0) + ep_size = num_experts # nle == 1 -> compact path + disp = self._make( + num_experts=num_experts, + top_k=top_k, + ep_size=ep_size, + sp_size=sp_size, + sp_rank=sp_rank, + ) + t_local = B * local_L + n = t_local * top_k + x_TD = torch.randn(t_local, D) + scores_TK = torch.rand(t_local, top_k) + ids_TK = torch.stack( + [torch.randperm(num_experts)[:top_k] for _ in range(t_local)] + ) + counts_E = _counts_from_ids(ids_TK, num_experts) + + mod = "torchtitan.models.common.token_dispatcher" + + def fake_dispatch_compact( + routed_ND, + packed_source_indices_N, + send_cnt_1E, + recv_capacity, + _ranks, + ): + recv = routed_ND.new_zeros((recv_capacity, routed_ND.shape[-1])) + counts = send_cnt_1E.clone() + offset = 0 + for d in range(ep_size): + for j in range(int(send_cnt_1E[0, d])): + recv[offset + j] = routed_ND[ + int(packed_source_indices_N[offset + j]) + ] + offset += int(send_cnt_1E[0, d]) + return recv, counts + + def fake_unpermute( + output_RD, + packed_source_indices_N, + send_cnt_1E, + _rc, + _ranks, + ): + assert packed_source_indices_N is not None + out = output_RD.new_zeros((n, output_RD.shape[-1])) + offset = 0 + for d in range(ep_size): + for j in range(int(send_cnt_1E[0, d])): + out[int(packed_source_indices_N[offset + j])] = output_RD[ + offset + j + ] + offset += int(send_cnt_1E[0, d]) + return out + + with ( + unittest.mock.patch( + f"{mod}.permute_a2av_dispatch_compact", fake_dispatch_compact + ), + unittest.mock.patch(f"{mod}.unpermute_a2av_combine_packed", fake_unpermute), + ): + routed_RD, _counts_e, md = disp.dispatch(x_TD, scores_TK, ids_TK, counts_E) + out = disp.combine( + routed_RD, + md, + x_TD, + num_local_tokens_after_padding=t_local, + local_seq_len_after_padding=local_L, + ) + # Identity experts -> each token's contribution is its score-weighted sum. + expected_rows = (x_TD.unsqueeze(1) * scores_TK.unsqueeze(-1)).sum(dim=1) + return disp, out, expected_rows, t_local + + def test_output_is_global_length(self): + """The SP buffer spans the WHOLE sequence, not just this rank's shard.""" + sp_size, local_L, B = 4, 3, 2 + _disp, out, _exp, t_local = self._run_ep_sp( + sp_size=sp_size, sp_rank=2, B=B, local_L=local_L + ) + self.assertEqual(out.shape[0], t_local * sp_size) + + def test_rows_land_at_global_offsets_and_elsewhere_zero(self): + """Each local token lands at b*global_seq + sp_rank*L + p; rest is zero.""" + sp_size, sp_rank, B, local_L = 4, 2, 2, 3 + disp, out, expected_rows, t_local = self._run_ep_sp( + sp_size=sp_size, sp_rank=sp_rank, B=B, local_L=local_L + ) + global_seq = local_L * sp_size + written = set() + for i in range(t_local): + g = (i // local_L) * global_seq + sp_rank * local_L + (i % local_L) + written.add(g) + torch.testing.assert_close(out[g], expected_rows[i], rtol=1e-4, atol=1e-4) + # Every row outside this rank's slice must be untouched (zero), so the + # downstream Partial -> Shard(seq) reduce cannot double-count. + for g in range(out.shape[0]): + if g not in written: + self.assertTrue( + torch.equal(out[g], torch.zeros_like(out[g])), + f"row {g} outside this rank's slice was written", + ) + + def test_sp1_is_local_length_and_unchanged(self): + """sp_size == 1 keeps the validated pure-EP behaviour: local-length output.""" + _disp, out, expected_rows, t_local = self._run_ep_sp( + sp_size=1, sp_rank=0, B=1, local_L=6 + ) + self.assertEqual(out.shape[0], t_local) + torch.testing.assert_close(out, expected_rows, rtol=1e-4, atol=1e-4) + + def test_distinct_sp_ranks_write_disjoint_slices(self): + """Two SP ranks of the same sequence never write the same global row.""" + sp_size, B, local_L = 4, 1, 4 + touched = [] + for sp_rank in (0, 3): + _d, out, _e, t_local = self._run_ep_sp( + sp_size=sp_size, sp_rank=sp_rank, B=B, local_L=local_L + ) + touched.append({g for g in range(out.shape[0]) if out[g].abs().sum() > 0}) + self.assertTrue(touched[0], "rank 0 wrote nothing") + self.assertTrue(touched[1], "rank 3 wrote nothing") + self.assertEqual(touched[0] & touched[1], set(), "SP slices overlap") + + +class TestPackedSourceIndices(unittest.TestCase): + def test_full_size_routing_avoids_sparse_tables(self): + generator = torch.Generator().manual_seed(42) + for ep in (64, 128): + with self.subTest(ep=ep): + dispatcher = _make_dispatcher(ep, 8, ep) + ids = torch.stack( + [torch.randperm(ep, generator=generator)[:8] for _ in range(1024)] + ) + flat_ids = ids.flatten() + with unittest.mock.patch( + "torch.nn.functional.one_hot", + side_effect=AssertionError( + "Packed routing must not build a sparse table" + ), + ): + counts, vector = dispatcher._build_packed_routing( + ids, _counts_from_ids(ids, ep) + ) + expected = torch.cat( + [torch.where(flat_ids == peer)[0] for peer in range(ep)] + ) + torch.testing.assert_close(vector.long(), expected, rtol=0, atol=0) + torch.testing.assert_close( + counts.flatten().long(), + torch.bincount(flat_ids, minlength=ep), + rtol=0, + atol=0, + ) + + def test_ep64_dispatch_metadata_contract(self): + ep_size, top_k, tokens, hidden = 64, 8, 17, 4 + dispatcher = _make_dispatcher(ep_size, top_k, ep_size) + n = tokens * top_k + inputs = torch.arange(tokens * hidden, dtype=torch.float32).reshape( + tokens, hidden + ) + ids = (torch.arange(n) % ep_size).reshape(tokens, top_k) + scores = torch.arange(n, dtype=torch.float32).reshape(tokens, top_k) + expected_counts = torch.bincount(ids.reshape(-1), minlength=ep_size).to( + torch.int32 + ) + expected_recv_counts = expected_counts.clone().reshape(1, ep_size) + expected_recv_counts[0, 1] += 1 + expected_recv_counts[0, 2] -= 1 + expected_vector = torch.tensor( + sorted(range(n), key=lambda row: int(ids.reshape(-1)[row])), + dtype=torch.int32, + ) + expected_capacity = dispatcher._packed_capacity(tokens, ep_size, 1) + + def fake_permute(routed, vector, counts, capacity, ranks): + torch.testing.assert_close(routed, inputs.repeat_interleave(top_k, dim=0)) + torch.testing.assert_close(counts.reshape(-1), expected_counts) + torch.testing.assert_close(vector, expected_vector) + self.assertEqual(capacity, expected_capacity) + self.assertEqual(ranks, [list(range(ep_size))]) + return routed.new_zeros((capacity, hidden)), expected_recv_counts + + with unittest.mock.patch( + "torchtitan.models.common.token_dispatcher.permute_a2av_dispatch_compact", + fake_permute, + ): + routed, local_counts, metadata = dispatcher.dispatch( + inputs, scores, ids, _counts_from_ids(ids, ep_size) + ) + + self.assertEqual(routed.shape, (expected_capacity, hidden)) + torch.testing.assert_close( + local_counts, torch.tensor([expected_capacity], dtype=torch.int32) + ) + self.assertEqual(metadata.compact_capacity, expected_capacity) + torch.testing.assert_close(metadata.packed_source_indices_N, expected_vector) + torch.testing.assert_close(metadata.send_counts_1E.reshape(-1), expected_counts) + torch.testing.assert_close(metadata.recv_counts_1E, expected_recv_counts) + torch.testing.assert_close( + metadata.token_indices_experts_sorted_N, torch.arange(n) // top_k + ) + torch.testing.assert_close( + metadata.topk_scores_experts_sorted_N, scores.reshape(-1) + ) + + def test_stable_destination_order_with_skew_and_zero_peers(self): + dispatcher = _make_dispatcher(8, 2, 8) + for ids in ( + torch.tensor([[7, 0], [7, 1], [0, 1], [7, 0]]), + torch.zeros((65, 2), dtype=torch.int64), + torch.arange(8).repeat(17).reshape(-1, 2), + ): + n = ids.numel() + counts, vector = dispatcher._build_packed_routing( + ids, _counts_from_ids(ids, 8) + ) + flat_ids = ids.reshape(-1) + expected = torch.cat( + [torch.where(flat_ids == peer)[0].to(torch.int32) for peer in range(8)] + ) + torch.testing.assert_close( + counts, _counts_from_ids(ids, 8).to(torch.int32).reshape(1, 8) + ) + self.assertTrue(torch.equal(vector, expected)) + self.assertEqual(int(counts.sum()), n) + self.assertTrue( + torch.equal(vector.sort().values, torch.arange(n, dtype=torch.int32)) + ) + + def test_distinct_topk_preserves_stable_order_and_peer_bounds(self): + tokens = 17 + generator = torch.Generator().manual_seed(918) + for ep, top_k in ((8, 1), (8, 2), (8, 8), (64, 8)): + dispatcher = _make_dispatcher(ep, top_k, ep) + patterns = { + "skew": torch.arange(ep - top_k, ep).repeat(tokens, 1), + "rotating": ( + torch.arange(tokens)[:, None] + torch.arange(top_k)[None, :] + ) + % ep, + "random": torch.stack( + [ + torch.randperm(ep, generator=generator)[:top_k] + for _ in range(tokens) + ] + ), + } + for pattern, ids in patterns.items(): + with self.subTest(ep=ep, top_k=top_k, pattern=pattern): + n = ids.numel() + counts, vector = dispatcher._build_packed_routing( + ids, _counts_from_ids(ids, ep) + ) + flat_ids = ids.reshape(-1).tolist() + expected = torch.tensor( + sorted(range(n), key=lambda row: flat_ids[row]), + dtype=torch.int32, + ) + torch.testing.assert_close(vector, expected, rtol=0, atol=0) + self.assertEqual(int(counts.sum()), n) + self.assertGreaterEqual(int(counts.min()), 0) + self.assertLessEqual(int(counts.max()), tokens) + torch.testing.assert_close( + vector.sort().values, + torch.arange(n, dtype=torch.int32), + rtol=0, + atol=0, + ) + expected_counts = torch.tensor( + [[flat_ids.count(peer) for peer in range(ep)]], + dtype=torch.int32, + ) + torch.testing.assert_close(counts, expected_counts, rtol=0, atol=0) + + def test_compile_dispatch_combine_with_nonuniform_gate_scores(self): + self._check_dispatch_combine_with_nonuniform_gate_scores(compiling=True) + + def test_eager_dispatch_combine_with_nonuniform_gate_scores(self): + self._check_dispatch_combine_with_nonuniform_gate_scores(compiling=False) + + def _check_dispatch_combine_with_nonuniform_gate_scores(self, compiling): + from torchtitan.distributed.neuron_ep import api + + ep, top_k, tokens, hidden = 8, 3, 5, 4 + dispatcher = _make_dispatcher(ep, top_k, ep) + ids = torch.tensor([[7, 1, 4], [0, 7, 1], [4, 6, 0], [2, 5, 7], [6, 2, 1]]) + inputs = ( + torch.arange(tokens * hidden, dtype=torch.float32) + .reshape(tokens, hidden) + .requires_grad_() + ) + scores = ( + torch.arange(1, tokens * top_k + 1, dtype=torch.float32).reshape( + tokens, top_k + ) + / 8 + ).requires_grad_() + seen_vectors = [] + + def expected_order(): + flat_ids = ids.reshape(-1).tolist() + return torch.tensor( + sorted(range(len(flat_ids)), key=lambda row: flat_ids[row]) + ) + + def permute(x, packed_source_indices_N, counts, capacity, ranks): + self.assertIsNotNone(packed_source_indices_N) + seen_vectors.append(packed_source_indices_N) + order = expected_order() + torch.testing.assert_close( + packed_source_indices_N.long(), order, rtol=0, atol=0 + ) + output = torch.cat([x[order], x.new_zeros(capacity - x.shape[0], hidden)]) + metadata = torch.zeros(4, ep, dtype=torch.int64) + metadata[2] = counts[0] * hidden + return output, metadata.to(torch.uint32) + + def unpermute(output, packed_source_indices_N, counts, recv_counts, ranks): + order = expected_order() + self.assertIsNotNone(packed_source_indices_N) + torch.testing.assert_close( + packed_source_indices_N.long(), order, rtol=0, atol=0 + ) + return output[: order.numel()][torch.argsort(order)] + + with ( + unittest.mock.patch("torch.compiler.is_compiling", return_value=compiling), + unittest.mock.patch.object(api, "_packed_permute_a2av_kernel", permute), + unittest.mock.patch.object(api, "_packed_unpermute_a2av_kernel", unpermute), + ): + packed, _, metadata = dispatcher.dispatch( + inputs, scores, ids, _counts_from_ids(ids, ep) + ) + self.assertIsNotNone(metadata.packed_source_indices_N) + # Distinct expert scales make wrong row ordering observable. + factors = (ids.reshape(-1) + 1).float()[ + metadata.packed_source_indices_N.long() + ] + factors = torch.cat( + [factors, factors.new_zeros(packed.shape[0] - factors.numel())] + ) + combined = dispatcher.combine( + packed * factors[:, None], + metadata, + inputs, + num_local_tokens_after_padding=tokens, + local_seq_len_after_padding=tokens, + ) + weights = ( + torch.arange(tokens * hidden, dtype=torch.float32).reshape( + tokens, hidden + ) + + 3 + ) + (combined * weights).sum().backward() + + scale = (scores.detach() * (ids + 1)).sum(1)[:, None] + torch.testing.assert_close(combined, inputs.detach() * scale, rtol=0, atol=0) + torch.testing.assert_close(inputs.grad, weights * scale, rtol=0, atol=0) + torch.testing.assert_close( + scores.grad, + (weights * inputs.detach()).sum(1)[:, None] * (ids + 1), + rtol=0, + atol=0, + ) + self.assertEqual(len(seen_vectors), 2) + self.assertTrue( + all(vector is metadata.packed_source_indices_N for vector in seen_vectors) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan_neuron/tests/test_ep_kernels.py b/torchtitan_neuron/tests/test_ep_kernels.py new file mode 100644 index 00000000000..4ae82e131e1 --- /dev/null +++ b/torchtitan_neuron/tests/test_ep_kernels.py @@ -0,0 +1,223 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Unit tests for the Neuron EP kernel wrappers and their tensor contracts.""" + +import unittest +import unittest.mock + +import torch + +from torchtitan.distributed.neuron_ep import ( + permute_a2av_dispatch_compact, + unpermute_a2av_combine_packed, +) + + +class TestPackedA2AVAutograd(unittest.TestCase): + @staticmethod + def _permute(hidden, indices, counts, capacity, _ranks): + output = hidden.new_zeros(capacity, hidden.shape[1]) + output[: hidden.shape[0]] = hidden[indices.long()] + metadata = torch.zeros(4, counts.shape[1], dtype=torch.uint32) + metadata[2] = counts[0].to(torch.uint32) * hidden.shape[1] + return output, metadata + + @staticmethod + def _unpermute( + output, + indices, + _send_counts, + _recv_counts, + _ranks, + ): + return output[: indices.shape[0]][torch.argsort(indices.long())] + + def test_kernel_io_rows_use_routed_and_expert_bounds(self): + local_tokens = 6 + top_k = 4 + # Match the dispatcher's supported one-expert-per-EP-rank capacity. + ep_size = 8 + hidden = 3 + routed_rows = local_tokens * top_k + expert_rows = local_tokens * ep_size + indices = torch.arange(routed_rows, dtype=torch.int32) + counts = torch.full((1, ep_size), routed_rows // ep_size, dtype=torch.int32) + api_globals = permute_a2av_dispatch_compact.__globals__ + + with unittest.mock.patch.dict( + api_globals, + { + "_packed_permute_a2av_kernel": self._permute, + "_packed_unpermute_a2av_kernel": self._unpermute, + }, + ): + routed = torch.randn(routed_rows, hidden, requires_grad=True) + packed, recv_counts = permute_a2av_dispatch_compact( + routed, indices, counts, expert_rows, [list(range(ep_size))] + ) + self.assertEqual(packed.shape, (expert_rows, hidden)) + packed.sum().backward() + self.assertEqual(routed.grad.shape, (routed_rows, hidden)) + + expert_output = torch.randn(expert_rows, hidden, requires_grad=True) + combined = unpermute_a2av_combine_packed( + expert_output, + indices, + counts, + recv_counts, + [list(range(ep_size))], + ) + self.assertEqual(combined.shape, (routed_rows, hidden)) + combined.sum().backward() + self.assertEqual(expert_output.grad.shape, (expert_rows, hidden)) + + def test_packed_dispatch_and_combine_are_adjoints(self): + n, hidden, capacity, ep_size = 8, 3, 16, 8 + indices = torch.arange(n, dtype=torch.int32) + counts = torch.ones(1, ep_size, dtype=torch.int32) + api_globals = permute_a2av_dispatch_compact.__globals__ + with unittest.mock.patch.dict( + api_globals, + { + "_packed_permute_a2av_kernel": self._permute, + "_packed_unpermute_a2av_kernel": self._unpermute, + }, + ): + routed = torch.randn(n, hidden, requires_grad=True) + packed, recv_counts = permute_a2av_dispatch_compact( + routed, indices, counts, capacity, [list(range(ep_size))] + ) + self.assertEqual(packed.shape, (capacity, hidden)) + packed.sum().backward() + torch.testing.assert_close(routed.grad, torch.ones_like(routed)) + + expert_output = torch.randn(capacity, hidden, requires_grad=True) + combined = unpermute_a2av_combine_packed( + expert_output, + indices, + counts, + recv_counts, + [list(range(ep_size))], + ) + combined.sum().backward() + expected = torch.zeros_like(expert_output) + expected[:n] = 1 + torch.testing.assert_close(expert_output.grad, expected) + + def test_recv_counts_divide_before_int32_narrowing(self): + hidden = 4096 + indices = torch.zeros(1, dtype=torch.int32) + counts = torch.tensor([[1, 0, 0, 0, 0, 0, 0, 0]], dtype=torch.int32) + + def permute_with_large_metadata(input_tensor, *_args): + metadata = torch.zeros(4, 8, dtype=torch.uint32) + metadata[2, 0] = 2**31 + return input_tensor, metadata + + api_globals = permute_a2av_dispatch_compact.__globals__ + with unittest.mock.patch.dict( + api_globals, + {"_packed_permute_a2av_kernel": permute_with_large_metadata}, + ): + _packed, recv_counts = permute_a2av_dispatch_compact( + torch.zeros(1, hidden), + indices, + counts, + 1, + [list(range(8))], + ) + self.assertEqual(int(recv_counts[0, 0]), 2**31 // hidden) + + +class TestPackedA2AVLNCContract(unittest.TestCase): + def test_lnc2_grid(self): + from torchtitan.distributed.neuron_ep.api import _packed_lnc2_grid + + self.assertEqual(_packed_lnc2_grid(2), (2,)) + + def test_non_lnc2_fails_before_kernel_launch(self): + from torchtitan.distributed.neuron_ep.api import _packed_lnc2_grid + + with self.assertRaisesRegex(RuntimeError, "NEURON_LOGICAL_NC_CONFIG=2"): + _packed_lnc2_grid(1) + + def test_replica_groups_require_complete_aligned_devices(self): + from torchtitan.distributed.neuron_ep.api import _validate_packed_ep_ranks + + _validate_packed_ep_ranks([list(range(8)), list(range(8, 16))]) + for invalid in ( + [list(range(4))], + [list(range(10))], + [list(range(2, 10))], + [[0, 1, 2, 3, 5, 6, 7, 8]], + ): + with ( + self.subTest(invalid=invalid), + self.assertRaisesRegex(ValueError, "packed A2AV"), + ): + _validate_packed_ep_ranks(invalid) + + +class TestPackedSourceIndexAutograd(unittest.TestCase): + def test_nonidentity_permutation_with_nonuniform_gradients(self): + n, h, c, ep = 7, 3, 13, 8 + vector = torch.tensor([5, 1, 6, 4, 0, 3, 2], dtype=torch.int32) + counts = torch.tensor([[2, 0, 1, 0, 3, 0, 0, 1]], dtype=torch.int32) + seen = [] + + def permute(x, packed_source_indices_N, cnt, cap, ranks): + seen.append(packed_source_indices_N) + torch.testing.assert_close(packed_source_indices_N, vector, rtol=0, atol=0) + output = torch.cat( + [x[packed_source_indices_N.long()], x.new_zeros(cap - n, h)] + ) + metadata = torch.zeros(4, ep, dtype=torch.int64) + metadata[2] = cnt[0] * h + return output, metadata.to(torch.uint32) + + def unpermute(output, packed_source_indices_N, cnt, recv_counts, ranks): + self.assertIs(packed_source_indices_N, vector) + order = vector.long() + return output[:n][torch.argsort(order)] + + with unittest.mock.patch.dict( + permute_a2av_dispatch_compact.__globals__, + { + "_packed_permute_a2av_kernel": permute, + "_packed_unpermute_a2av_kernel": unpermute, + }, + ): + ranks = [list(range(ep))] + routed = ( + torch.arange(n * h, dtype=torch.float32).reshape(n, h).requires_grad_() + ) + packed, recv_counts = permute_a2av_dispatch_compact( + routed, vector, counts, c, ranks + ) + packed_grad = torch.arange(c * h, dtype=torch.float32).reshape(c, h) + 2 + packed.backward(packed_grad) + torch.testing.assert_close( + routed.grad, + packed_grad[:n][torch.argsort(vector.long())], + rtol=0, + atol=0, + ) + expert = ( + torch.arange(c * h, dtype=torch.float32).reshape(c, h).requires_grad_() + ) + combined = unpermute_a2av_combine_packed( + expert, vector, counts, recv_counts, ranks + ) + combined_grad = torch.arange(n * h, dtype=torch.float32).reshape(n, h) + 11 + combined.backward(combined_grad) + self.assertEqual(len(seen), 2) + self.assertTrue(all(x is vector for x in seen)) + expected = torch.cat([combined_grad[vector.long()], torch.zeros(c - n, h)]) + torch.testing.assert_close(expert.grad, expected, rtol=0, atol=0) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan_neuron/tests/test_flex_attention.py b/torchtitan_neuron/tests/test_flex_attention.py new file mode 100644 index 00000000000..ee27a8d04ab --- /dev/null +++ b/torchtitan_neuron/tests/test_flex_attention.py @@ -0,0 +1,139 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Neuron FlexAttention kernel-lowering tests. + +Exercise static causal and sliding-window masks and dynamic document masks +directly, without modifying the framework's document-mask builders. +""" + +import logging +import unittest + +import pytest +import torch + +pytest.importorskip("torch_neuronx") + +from torchtitan.tools.utils import device_type # noqa: E402 + + +_NKI_LOGGER = ( + "torch_neuronx.neuron_dynamo_backend.fx.passes.flex_attention_nki_lowering" +) +_FALLBACK_MARKERS = ( + "falling back to graph-based", + "Failed to lower flex_attention", + "not statically computable", +) + + +class _FallbackCatcher(logging.Handler): + """Record whether the NKI-lowering pass fell back for flex.""" + + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.records: list[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + if record.name.startswith(_NKI_LOGGER) and any( + m in record.getMessage() for m in _FALLBACK_MARKERS + ): + self.records.append(record.getMessage()) + + +@unittest.skipUnless(device_type == "neuron", "requires a Neuron device") +class TestFlexNKILoweringOnNeuron(unittest.TestCase): + """Check NKI lowering for static masks and fallback for dynamic masks.""" + + B, H, S, D = 1, 4, 512, 64 + + def setUp(self): + # Isolate torch.compile / dynamo state per test: the NKI-lowering pass + # caches per compiled graph, so a fallback logged while compiling one + # mask must not leak into a later test that compiles a different mask. + torch._dynamo.reset() + + def tearDown(self): + torch._dynamo.reset() + + def _run_flex(self, mask_mod, seq_len: int | None = None) -> bool: + """Compile + run fwd/bwd for one mask_mod; True if it lowered to NKI.""" + from torch.nn.attention.flex_attention import create_block_mask, flex_attention + + s = self.S if seq_len is None else seq_len + catcher = _FallbackCatcher() + root = logging.getLogger() + root.addHandler(catcher) + try: + compiled_cbm = torch.compile(create_block_mask, backend="neuron") + compiled_flex = torch.compile(flex_attention, backend="neuron") + + torch.manual_seed(0) + qkv = [ + torch.randn( + self.B, + self.H, + s, + self.D, + device="neuron", + dtype=torch.bfloat16, + requires_grad=True, + ) + for _ in range(3) + ] + block_mask = compiled_cbm(mask_mod, self.B, self.H, s, s, device="neuron") + out = compiled_flex(*qkv, block_mask=block_mask) + out.sum().backward() + finally: + root.removeHandler(catcher) + return len(catcher.records) == 0 + + def test_static_causal_mask_lowers_to_nki(self): + def causal(b, h, q_idx, kv_idx): + return q_idx >= kv_idx + + self.assertTrue( + self._run_flex(causal), + "static causal mask_mod should lower to the NKI kernel (no fallback)", + ) + + def test_static_sliding_window_mask_lowers_to_nki(self): + from torch.nn.attention.flex_attention import and_masks + + from torchtitan.models.common.attention import ( + get_causal_mask_mod, + get_sliding_window_mask_mod, + ) + + seq_len = 1024 + composed = and_masks( + get_causal_mask_mod(), + get_sliding_window_mask_mod(128), + ) + + self.assertTrue( + self._run_flex(composed, seq_len=seq_len), + "static causal sliding-window mask should lower to the NKI kernel " + "(no fallback)", + ) + + def test_dynamic_document_mask_falls_back(self): + # A mask_mod that indexes a captured tensor is dynamic -> NKI cannot + # statically evaluate it. + doc_ids = torch.zeros(self.S, dtype=torch.int32, device="neuron") + doc_ids[self.S // 2 :] = 1 + + def document(b, h, q_idx, kv_idx): + return (q_idx >= kv_idx) & (doc_ids[q_idx] == doc_ids[kv_idx]) + + self.assertFalse( + self._run_flex(document), + "dynamic (captured-tensor) mask_mod is expected to fall back from NKI", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan_neuron/tests/test_metrics_mfu.py b/torchtitan_neuron/tests/test_metrics_mfu.py new file mode 100644 index 00000000000..a51cefe07c4 --- /dev/null +++ b/torchtitan_neuron/tests/test_metrics_mfu.py @@ -0,0 +1,49 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""Unit tests for MFU computation, including the Neuron MXFP8 4x-peak path. + +These test the extracted ``compute_mfu`` helper directly (pure math), so they +run on CPU without real Neuron hardware: the Neuron branch is exercised by +patching ``has_neuron_capability`` rather than requiring a trn3 device. +""" + +from unittest import mock + +from torchtitan.components import metrics +from torchtitan.components.metrics import compute_mfu + +# Simple round numbers so the expected values are obvious. +_FLOPS_PER_TOKEN = 100 +_TPS = 10.0 +_PEAK = 1000.0 + + +def test_unquantized_mfu_uses_bf16_peak(): + # 100 * 100 * 10 / 1000 = 100.0 + mfu = compute_mfu(_FLOPS_PER_TOKEN, _TPS, _PEAK, has_quantization=False) + assert mfu == 100 * _FLOPS_PER_TOKEN * _TPS / _PEAK + + +def test_quantized_non_neuron_mfu_is_none(): + # Off Neuron, a quantized run cannot report a meaningful MFU. + with mock.patch.object(metrics.utils, "has_neuron_capability", return_value=False): + mfu = compute_mfu(_FLOPS_PER_TOKEN, _TPS, _PEAK, has_quantization=True) + assert mfu is None + + +def test_quantized_neuron_mxfp8_uses_4x_peak(): + # On trn3 MXFP8, peak is 4x bf16, so MFU is 1/4 of the naive value. + with mock.patch.object(metrics.utils, "has_neuron_capability", return_value=True): + mfu = compute_mfu(_FLOPS_PER_TOKEN, _TPS, _PEAK, has_quantization=True) + assert mfu == 100 * _FLOPS_PER_TOKEN * _TPS / (4 * _PEAK) + + +def test_neuron_mxfp8_mfu_is_quarter_of_unquantized(): + # Cross-check the 4x relationship between the two branches directly. + unquant = compute_mfu(_FLOPS_PER_TOKEN, _TPS, _PEAK, has_quantization=False) + with mock.patch.object(metrics.utils, "has_neuron_capability", return_value=True): + neuron = compute_mfu(_FLOPS_PER_TOKEN, _TPS, _PEAK, has_quantization=True) + assert neuron == unquant / 4 diff --git a/torchtitan_neuron/tests/test_mxfp8_numerics.py b/torchtitan_neuron/tests/test_mxfp8_numerics.py new file mode 100644 index 00000000000..81b68c387d9 --- /dev/null +++ b/torchtitan_neuron/tests/test_mxfp8_numerics.py @@ -0,0 +1,238 @@ +# Copyright (c) Amazon.com, Inc. or its affiliates +# All rights reserved. +# +# Licensed under the BSD-3-Clause license. + +"""MXFP8-vs-bf16 accuracy tests for the titan MXFP8 linear on Neuron trn3. + +Builds the torchtitan ``MXFP8Linear`` (which wraps torchao's ``MXFP8Linear`` +and, on trn3, selects the ``neuron`` GEMM kernel), runs a forward ++ backward on device, and checks the output and both gradients track a bf16 stock +reference within an MXFP8 quantization-error bar (eager: the bare linear vs +``F.linear``). + +Gated on a real trn3 device: the ``neuron`` kernel + backend only dispatch there. +Shapes are kernel-valid (features multiples of 128; tokens multiple of 128). +""" + +import unittest + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import distribute_tensor, Replicate, Shard +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) + +from torchtitan.tools.utils import has_neuron_capability + +from torchtitan_neuron.components.quantization import torchao_supports_neuron + +pytestmark = pytest.mark.skipif( + not torchao_supports_neuron(), + reason="requires torchao with Neuron MXFP8 support", +) + +# Kernel-valid shapes: in/out features multiples of 128, tokens multiple of 128. +_IN, _OUT, _TOKENS = 512, 512, 128 + +# MXFP8-vs-bf16 relative-error bars. Measured on trn3: eager fwd ~0.075, +# dgrad ~0.052, wgrad ~0.057. Bars carry headroom for seed variation. +_FWD_BAR = 0.10 +_EAGER_GRAD_BAR = 0.10 + + +def _rel_error(actual: torch.Tensor, ref: torch.Tensor) -> float: + """Relative L2 error ``||actual - ref|| / ||ref||`` (both as f32 on CPU).""" + a = actual.detach().to("cpu").float() + r = ref.detach().to("cpu").float() + return (torch.linalg.norm(a - r) / torch.linalg.norm(r)).item() + + +def _build_mxfp8_linear() -> tuple["nn.Module", torch.Tensor]: + """Build a Neuron MXFP8 linear with a known bf16 weight (the reference ground + truth). Cast to bf16 up front: the Neuron MXFP8 GEMM takes bf16 operands (an + fp32 module, the build default, is rejected by the fused op).""" + from torchtitan.components.quantization.mx import MXFP8Linear + + lin = MXFP8Linear.Config( + in_features=_IN, + out_features=_OUT, + bias=False, + kernel_preference="neuron", + ).build() + weight = torch.randn(_OUT, _IN, dtype=torch.bfloat16) + with torch.no_grad(): + lin.weight.copy_(weight) + return lin, weight + + +@unittest.skipUnless( + has_neuron_capability("mxfp8"), "requires a real Neuron trn3 device" +) +class TestMXFP8LinearNumerics(unittest.TestCase): + def _run_numerics(self, *, module, linear, weight, ref_fn, compiled, grad_bar): + torch.manual_seed(0) + runner = ( + torch.compile(module, backend="neuron", fullgraph=True) + if compiled + else module + ) + + x_cpu = torch.randn(_TOKENS, _IN, dtype=torch.bfloat16) + x_dev = x_cpu.detach().to("neuron").requires_grad_(True) + out = runner(x_dev) + self.assertEqual(tuple(out.shape), (_TOKENS, _OUT)) + self.assertTrue(torch.isfinite(out.to("cpu")).all()) + out.sum().backward() + + # Gradients must propagate to the input and the weight (finite, right shape). + self.assertIsNotNone(x_dev.grad) + self.assertIsNotNone(linear.weight.grad) + self.assertEqual(tuple(linear.weight.grad.shape), (_OUT, _IN)) + self.assertTrue(torch.isfinite(x_dev.grad.to("cpu")).all()) + self.assertTrue(torch.isfinite(linear.weight.grad.to("cpu")).all()) + + # bf16 float reference from the same bf16-rounded weight. + x_ref = x_cpu.detach().float().requires_grad_(True) + w_ref = weight.detach().float().requires_grad_(True) + out_ref = ref_fn(x_ref, w_ref) + out_ref.sum().backward() + + fwd_rel = _rel_error(out, out_ref) + dgrad_rel = _rel_error(x_dev.grad, x_ref.grad) + wgrad_rel = _rel_error(linear.weight.grad, w_ref.grad) + self.assertLess(fwd_rel, _FWD_BAR, f"fwd rel error {fwd_rel:.4f}") + self.assertLess(dgrad_rel, grad_bar, f"input-grad rel error {dgrad_rel:.4f}") + self.assertLess(wgrad_rel, grad_bar, f"weight-grad rel error {wgrad_rel:.4f}") + + def test_forward_backward_vs_bf16_linear(self): + lin, weight = _build_mxfp8_linear() + lin = lin.to(torch.bfloat16).to("neuron") + self._run_numerics( + module=lin, + linear=lin, + weight=weight, + ref_fn=F.linear, + compiled=False, + grad_bar=_EAGER_GRAD_BAR, + ) + + +@unittest.skip( + "DTensor + compile spawns worker processes (MultiProcessTestCase); the Neuron " + "runtime cannot re-initialize in the children after the parent used it in a " + "shared test session (Failed to initialize Neuron Runtime: status code 1). " + "The DTensor + compile + NKI path is covered by the qwen3_8b_mxfp8 (tp4 x " + "fsdp2) E2E compile run instead." +) +@unittest.skipUnless( + has_neuron_capability("mxfp8"), "requires a real Neuron trn3 device" +) +class TestMXFP8LinearCompileDTensor(DTensorTestBase): + """MXFP8 linear with a DTensor (TP-sharded) weight under torch.compile.""" + + @property + def world_size(self) -> int: + return 2 + + @with_comms + def test_tp_sharded_linear_compiles_and_runs(self): + from torchtitan.components.quantization.mx import MXFP8Linear + + torch.manual_seed(0) + mesh = init_device_mesh( + self.device_type, (self.world_size,), mesh_dim_names=("tp",) + ) + + lin = ( + MXFP8Linear.Config( + in_features=_IN, + out_features=_OUT, + bias=False, + kernel_preference="neuron", + ) + .build() + .to(torch.bfloat16) + .to(self.device_type) + ) + + # Column-parallel linear: shard the weight rows (out_features) across TP, + # replicate the input. Output is sharded on the last dim. + with torch.no_grad(): + w_dt = distribute_tensor(lin.weight.detach(), mesh, [Shard(0)]) + lin.weight = nn.Parameter(w_dt) + + x = distribute_tensor( + torch.randn(_TOKENS, _IN, dtype=torch.bfloat16, device=self.device_type), + mesh, + [Replicate()], + ).requires_grad_(True) + + compiled = torch.compile(lin, backend="neuron", fullgraph=True) + out = compiled(x) + + full = out.full_tensor() + self.assertEqual(tuple(full.shape), (_TOKENS, _OUT)) + self.assertTrue(torch.isfinite(full.to("cpu")).all()) + + out.sum().backward() + self.assertIsNotNone(lin.weight.grad) + wgrad_full = lin.weight.grad.full_tensor() + self.assertEqual(tuple(wgrad_full.shape), (_OUT, _IN)) + self.assertTrue(torch.isfinite(wgrad_full.to("cpu")).all()) + + +# --- Neuron trn3 MXFP8 capability gating + kernel_preference flow --- + + +def test_mxfp8_linear_converter_accepts_neuron_trn3(monkeypatch): + """MXFP8LinearConverter constructs on Neuron trn3 (no CUDA SM100).""" + pytest.importorskip("torchao") + from torchtitan.components.quantization import mx + + monkeypatch.setattr(mx, "has_cuda_capability", lambda *a: False) + monkeypatch.setattr(mx, "has_neuron_capability", lambda feature: True) + # Should not raise: Neuron trn3 satisfies the capability gate. + mx.MXFP8LinearConverter(mx.MXFP8LinearConverter.Config(model_compile_enabled=False)) + + +def test_mxfp8_linear_converter_rejects_unsupported_device(monkeypatch): + """Neither CUDA SM100 nor Neuron trn3 -> ValueError naming trn3.""" + pytest.importorskip("torchao") + from torchtitan.components.quantization import mx + + monkeypatch.setattr(mx, "has_cuda_capability", lambda *a: False) + monkeypatch.setattr(mx, "has_neuron_capability", lambda feature: False) + with pytest.raises(ValueError, match="trn3"): + mx.MXFP8LinearConverter( + mx.MXFP8LinearConverter.Config(model_compile_enabled=False) + ) + + +def test_kernel_preference_flows_to_torchao(): + """An explicit kernel_preference="neuron" reaches the torchao MXFP8Linear. + + Skips unless the installed torchao exposes KernelPreference.NEURON; + upstream torchao has no NEURON member. + """ + pytest.importorskip("torchao") + kp = pytest.importorskip("torchao.quantization.quantize_.common.kernel_preference") + if not hasattr(kp.KernelPreference, "NEURON"): + pytest.skip("installed torchao has no KernelPreference.NEURON") + + from torchtitan.components.quantization.mx import MXFP8Linear + + module = MXFP8Linear.Config( + in_features=16, out_features=32, bias=False, kernel_preference="neuron" + ).build() + assert module.kernel_preference == kp.KernelPreference.NEURON + + +if __name__ == "__main__": + unittest.main()