diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 3ca7754175b..d939d68729c 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,3 +1,6 @@ +# Files directly in the repository root without a more specific owner. +/* @NVIDIA/core-nemo + megatron/core/ @NVIDIA/core-adlr @NVIDIA/core-nemo megatron/core/tensor_parallel/generalized_tensor_parallelism.py @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/gtp diff --git a/.github/actions/check-nvidia-sso-membership/action.yml b/.github/actions/check-nvidia-sso-membership/action.yml index 71926c4547d..e96863e2264 100644 --- a/.github/actions/check-nvidia-sso-membership/action.yml +++ b/.github/actions/check-nvidia-sso-membership/action.yml @@ -1,5 +1,5 @@ name: 'Check NVIDIA SSO Membership' -description: 'Check if a GitHub username exists in the NVIDIA SSO users list from github-audits' +description: 'Check if a GitHub username is in the NVIDIA SSO users list or internal allowlist' author: 'NVIDIA' inputs: @@ -24,7 +24,7 @@ inputs: outputs: is_member: - description: 'Boolean - true if user is in NVIDIA SSO list, false otherwise' + description: 'Boolean - true if user is in NVIDIA SSO list or internal allowlist, false otherwise' value: ${{ steps.check-membership.outputs.is_member }} is_org_member: description: 'Boolean - true if user has NVIDIA or NVIDIA-NeMo in org_roles' @@ -91,6 +91,15 @@ runs: echo "Checking if $USERNAME is in NVIDIA SSO users list..." + # Service accounts cannot enroll in SSO but can be explicitly trusted. + if [ "$USERNAME" = "svcnemo-autobot" ]; then + echo "$USERNAME found in NVIDIA internal user allowlist" + echo "is_member=true" >> $GITHUB_OUTPUT + echo "is_org_member=false" >> $GITHUB_OUTPUT + echo "user_orgs=" >> $GITHUB_OUTPUT + exit 0 + fi + # Check if SSO file is available if [ "${{ steps.download-sso.outputs.sso_file_available }}" != "true" ] || [ ! -f "$SSO_FILE" ]; then echo "ERROR: $SSO_FILE not available - cannot check membership" diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml index eface2e458a..9b33b89739d 100644 --- a/.github/copy-pr-bot.yaml +++ b/.github/copy-pr-bot.yaml @@ -1,4 +1,4 @@ enabled: true auto_sync_draft: false auto_sync_ready: true -trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "CarlosGomes98", "ChenhanYu", "Connor-XY", "DanialTaheri", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JF-D", "JRD971000", "Leili", "Mellonta", "Phlip79", "QiZhangNV", "RPrenger", "ShriyaRishab", "WanZzzzzz", "Wohox", "YangFei1990", "ZhiyuLi-Nvidia", "adistomar", "ahmadki", "aklife97", "alokpathy", "ananthsub", "anlthms", "aroshanghias-nvd", "ashehper", "asolergi-nv", "athitten", "balasaajay", "buptzyb", "chtruong814", "cjld", "cspades", "cuichenx", "deepakn94", "desh2608", "dimapihtar", "dingqingy-nv", "duncanriach", "ehosseiniasl", "erhoo82", "ericharper", "fanshiqing", "faradawn", "fitsumreda", "freewym", "frsun-nvda", "gautham-kollu", "gdengk", "goelarushi", "guihong-nv", "guyueh1", "hexinw-nvidia", "huvunvidia", "hxbai", "ilml", "jalbericiola", "janEbert", "jaredcasper", "jenchen13", "jiaji-huang", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kajalj22", "kamran-nvidia", "kevalmorabia97", "kevjshih", "kingformatty", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwyss-nvidia", "lauradang", "layalir", "lhb8125", "liding-nv", "lmcafee-nvidia", "maanug-nv", "macandro96", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "minitu", "mkhona-nvidia", "nanz-nv", "ntajbakhsh", "nvcsathe", "parthmannan", "philipcmonk", "prajwal1210", "pthombre", "rapatel", "rhewett-nv", "rogerwaleffe", "sajadn", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "sharathts", "sheliang-nv", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sraman-rgb", "sudhakarsingh27", "svcnvidia-nemo-ci", "tdene", "theothermike", "thomasdhc", "tomlifu", "trintamaki", "tylerpoon", "wdykas", "wujingyue", "xiaoyao0115", "xuantengh", "xuwchen", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yqwangustc", "yueshen2016", "yuzhongw-nvidia", "zhehuaichen", "zhongbozhu"] +trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "CarlosGomes98", "ChenhanYu", "Connor-XY", "DanialTaheri", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JF-D", "JRD971000", "Leili", "Mellonta", "Phlip79", "QiZhangNV", "RPrenger", "ShriyaRishab", "WanZzzzzz", "Wohox", "YangFei1990", "ZhiyuLi-Nvidia", "adistomar", "ahmadki", "aklife97", "alokpathy", "ananthsub", "anlthms", "aroshanghias-nvd", "ashehper", "asolergi-nv", "athitten", "balasaajay", "buptzyb", "chtruong814", "cjld", "cspades", "cuichenx", "deepakn94", "desh2608", "dimapihtar", "dingqingy-nv", "duncanriach", "ehosseiniasl", "erhoo82", "ericharper", "fanshiqing", "faradawn", "fitsumreda", "freewym", "frsun-nvda", "gautham-kollu", "gdengk", "goelarushi", "guihong-nv", "guyueh1", "hexinw-nvidia", "huvunvidia", "hxbai", "ilml", "jalbericiola", "janEbert", "jaredcasper", "jenchen13", "jiaji-huang", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kajalj22", "kamran-nvidia", "kevalmorabia97", "kevjshih", "kingformatty", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwyss-nvidia", "lauradang", "layalir", "lhb8125", "liding-nv", "lmcafee-nvidia", "maanug-nv", "macandro96", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "minitu", "mkhona-nvidia", "nanz-nv", "niyunsheng", "ntajbakhsh", "nvcsathe", "parthmannan", "philipcmonk", "prajwal1210", "pthombre", "rapatel", "rhewett-nv", "rogerwaleffe", "sajadn", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "sharathts", "sheliang-nv", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sraman-rgb", "sudhakarsingh27", "svcnvidia-nemo-ci", "tdene", "theothermike", "thomasdhc", "tomlifu", "trintamaki", "tylerpoon", "vasunvidia", "wdykas", "wujingyue", "xiaoyao0115", "xuantengh", "xuwchen", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yqwangustc", "yueshen2016", "yuzhongw-nvidia", "zhehuaichen", "zhongbozhu"] diff --git a/.github/oncall_schedule.json b/.github/oncall_schedule.json index 0fe3ecb86ab..622a830a73e 100644 --- a/.github/oncall_schedule.json +++ b/.github/oncall_schedule.json @@ -1,8 +1,4 @@ [ - { - "user": "janEbert", - "date": "2026-08-05" - }, { "user": "maanug-nv", "date": "2026-08-12" @@ -46,5 +42,9 @@ { "user": "janEbert", "date": "2026-10-21" + }, + { + "user": "maanug-nv", + "date": "2026-10-28" } ] diff --git a/.gitlab/stages/06.triage.yml b/.gitlab/stages/06.triage.yml index 91f8a4e703c..bac492bc873 100644 --- a/.gitlab/stages/06.triage.yml +++ b/.gitlab/stages/06.triage.yml @@ -49,11 +49,17 @@ triage:linear_write: artifacts: true allow_failure: true script: - - >- - cerno-linear write - --config "${CERNO_CONFIG}" - --plan linear_action_plan.json - --output linear_action_plan_post.json + # The shared GitHub email lookup reads the existing GH_TOKEN CI variable. + - | + ONCALL_ASSIGNEE="$( + python tests/test_utils/python_scripts/resolve_oncall_assignee.py \ + --schedule-file .github/oncall_schedule.json + )" + cerno-linear write \ + --config "${CERNO_CONFIG}" \ + --plan linear_action_plan.json \ + --output linear_action_plan_post.json \ + --assignee "${ONCALL_ASSIGNEE}" artifacts: when: always paths: diff --git a/docker/Dockerfile.linting b/docker/Dockerfile.linting index e9a4de80455..ef3a5813e6a 100644 --- a/docker/Dockerfile.linting +++ b/docker/Dockerfile.linting @@ -24,7 +24,7 @@ RUN --mount=type=secret,id=JET_INDEX_URLS \ # Keep this in the internal-only stage so public CI has no internal service dependency. ARG CI_SERVER_URL -ARG CERNO_COMMIT=5a5fb5360e67f8f09d189871bbc0d768c09c43fa +ARG CERNO_COMMIT=3219afcbfcdbddd1d55e013d94e54a15986e2b1f RUN --mount=type=secret,id=CERNO_TOKEN \ GIT_CONFIG_COUNT=1 \ GIT_CONFIG_KEY_0=http.extraHeader \ diff --git a/docs/api-guide/core/generalized_tensor_parallel.md b/docs/api-guide/core/generalized_tensor_parallel.md index 831db034d90..11c01b8a504 100644 --- a/docs/api-guide/core/generalized_tensor_parallel.md +++ b/docs/api-guide/core/generalized_tensor_parallel.md @@ -417,7 +417,7 @@ The figure visualizes the per-class split from the list above: green = resolves Two distinct pools with explicit lifecycle rules: -- **`GTPWeightCache`** (AG/RS output buffers) — ticket-based, keyed on `(shape, dtype, fwd, expert_idx, reduce_scatter)`. Same-shape buffers across layers are shared, **except between chain neighbours** — one-step-ahead keeps `prev_w` and the current weight live at once, so `_ensure_distinct_buffer_from_prev` folds a parity bit into the key when the two would collide, at the cost of one extra buffer for the second of the pair. Normally inert (neighbours are different weight roles, hence different shapes); it fires when CG capture leaves two same-shaped weights adjacent — embedding + output_layer alone in the `UNGRAPHED` chain. Tickets persistent; buffer allocated lazily on first `get()`; addresses stable across iterations for CG replay. +- **`GTPWeightCache`** (AG/RS output buffers) — ticket-based, keyed on `(shape, dtype, fwd, expert_idx, reduce_scatter)`, plus a `("recompute", parity)` suffix for recompute-chain gathers. Same-shape buffers across layers are shared, **except between chain neighbours** — one-step-ahead keeps the predecessor and the current weight live at once, so `_ensure_no_shared_buffer_with` folds a parity bit into the key when the two would collide, at the cost of one extra buffer for the second of the pair. The caller names which chain to guard, because the chains disagree on who a weight's neighbour is: on the fwd chain the check is normally inert (neighbours are different roles, hence different shapes) and fires only when CG capture leaves two same-shaped weights adjacent — embedding + output_layer alone in the `UNGRAPHED` chain — whereas on a recompute chain same-shape adjacency is the norm. Tickets persistent; buffer allocated lazily on first `get()`; addresses stable across iterations for CG replay. - **`_wgrad_buf_pool`** (wgrad-GEMM output recycling) — holds the **full, unsharded** wgrad-GEMM output buffer (shape `_unsharded_shape`, dtype `main_grad.dtype` — fp32 when `grad_reduce_in_fp32`, else bf16). The TE backward writes the wgrad into it via `main_grad_func = weight.grad_buffer` (a `DistributedWeight` protocol method backed by `get_wgrad_tensor`; it is a *scratch*, distinct from the sharded `param.main_grad`); the protocol's `finalize_group_grads` (backed by `wgrad_reduce_scatter`) then reduce-scatters it down to the shard and the buffer is returned here. This is a full-weight-shaped fp32/bf16 transient — one of the larger per-weight buffers — and is **precision-independent** (wgrad is always computed in high precision), so it is identical in BF16 vs MXFP8 runs. Buffers are tagged `_from_gtp_wgrad_pool=True` at `_wgrad_pool_get`; `_wgrad_pool_put` no-ops on foreign buffers (fresh allocs from Megatron `layers.py` or aten F.embedding bwd) → caching allocator handles those, so the pool never accumulates untagged buffers. #### Overlap design summary @@ -433,7 +433,7 @@ GTP_remat runs up to **three** independent prefetch chains, all following one ru |---|------|---------|--------------------|--------------|------| | 1 | fwd | weight `i` | `next_w` = i+1 ‖ `GEMM_i` | rowwise (`fwd=True`) | `_prefetch_handle` | | 2 | bwd dgrad | weight `i` | `prev_w` = i−1 ‖ `Dgrad_i` | columnwise (`fwd=False`) | `_prefetch_handle` | -| 3 | bwd recompute | weight `i` | `_recompute_next` = i+1 ‖ `recompute_GEMM_i` | rowwise (`fwd=True`) | `_recompute_prefetch_handle` (separate) | +| 3 | bwd recompute | weight `i` | `_recompute_next` = i+1 ‖ `recompute_GEMM_i` | rowwise (`fwd=True`) | `_recompute_prefetch_handle` + `_ag_ticket_recompute` (separate) | | 1b | fwd (MoE, eager) | expert weight `i` | same role in MoE block i+1 ‖ *whole block i* | rowwise (`fwd=True`) | `_prefetch_handle` | Row 1b is chain 1 applied to a *homogeneous* chain: routed-expert `fc1`/`fc2` link across consecutive MoE blocks, so the runway is a full block rather than one GEMM (§3.4 *Grouped-expert chains*). @@ -452,7 +452,7 @@ A future MR will add an opt-in wgrad-before-dgrad schedule on `_Linear` / `_Laye ##### Recompute-forward prefetch chain *(GTP_remat + activation recompute)* -When a GTP_remat-sharded module is in `--recompute-modules` (e.g. `shared_experts`), its forward is **re-run during backward** to regenerate activations. That recompute-forward must all-gather each weight **rowwise** again — a *third* gather lifecycle, concurrent with the in-flight **columnwise** dgrad gather of the *same* weight. Since both share one `GTPShardedParam`, the recompute path gets its **own** prefetch slot (`_recompute_prefetch_handle` / `_recompute_ag_event`, reusing the `_ag_ticket_fwd` rowwise buffer) so it never clobbers the dgrad lifecycle's `state` / `_prefetch_handle` / `ag_event`. +When a GTP_remat-sharded module is in `--recompute-modules` (e.g. `shared_experts`), its forward is **re-run during backward** to regenerate activations. That recompute-forward must all-gather each weight **rowwise** again — a *third* gather lifecycle, concurrent with the in-flight **columnwise** dgrad gather of the *same* weight. Since both share one `GTPShardedParam`, the recompute path gets its **own** prefetch slot (`_recompute_prefetch_handle` / `_recompute_ag_event`) so it never clobbers the dgrad lifecycle's `state` / `_prefetch_handle` / `ag_event`, and its **own** buffer ticket (`_ag_ticket_recompute`) with a parity of its own. Reusing `_ag_ticket_fwd` is unsafe twice over: a fwd prefetch may still be in flight in that buffer, and the fwd parity is decided against `prev_w` — a different neighbour. Without its own parity, consecutive recompute nodes share one buffer and the one-ahead prefetch overwrites the weight still being read: silent wrong activations, then NaN. The recompute weights form a **separate** linked list (`_recompute_next`), **self-populated** on the first backward from the weights actually re-gathered while `in_fp8_activation_recompute_phase()` is true — membership is *observed*, not configured (no tagging, so it tracks exactly what each checkpointed module re-gathers). Each recompute-forward consume prefetches the next recompute weight, so every gather **except the global-first** overlaps preceding recompute / dgrad / wgrad compute: @@ -606,7 +606,7 @@ Three consequences: - one-block-ahead makes block *N* and block *N+1* weights **live at the same time** — same key, two tensors in flight; - fix: a chain-position **parity (0,1,0,1…)** is folded into the cache key, so consecutive blocks alternate between **exactly two** buffers (counter cleared by `reset_gtp_state()`); - without it the prefetch would **overwrite the weight the running GEMM is still reading** — a silent-correctness bug, not a crash; - - the hazard is not exclusive to grouped chains — *any* chain whose neighbours share a key has it. Grouped chains are same-key throughout, so they take the blanket counter; others take the narrower `_ensure_distinct_buffer_from_prev` check, which allocates only where the collision is real (see *Buffer / memory management*). + - the hazard is not exclusive to grouped chains — *any* chain whose neighbours share a key has it. Grouped chains are same-key throughout, so they take the blanket counter; others take the narrower `_ensure_no_shared_buffer_with` check, which allocates only where the collision is real (see *Buffer / memory management*). - **Eager only** — the optimization disables itself under CUDA-graph capture: - `_classify_param_chain` evaluates `graphed = _FULL_ITERATION or ("moe" in cuda_graph_modules)` **before** the split, and returns the plain `GRAPHED` chain when it is true; - so with `--cuda-graph-impl full_iteration` **every** param is `GRAPHED` — expert weights included — and they keep the ordinary one-step-ahead prefetch; @@ -741,6 +741,7 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall | `test_gtp_partial_cg.py` | Four-layer partial-CG loss and eager-vs-replay grad-norm parity with two-slot ring reuse across independently replayed graphs (§3.5). | | `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. | | `test_gtp_muon_dcp.py` | Muon optimizer-state DCP roundtrip (§1.6): `replica_id` fold + native-FP8 backfill matching. | +| `test_gtp_recompute_chain.py` | Recompute-chain buffers (§3.1): adjacent nodes never share a gather buffer, dense and grouped, plus dgrad/wgrad parity vs no-recompute. | | `test_gtp_mtp.py` | GTP_remat + MTP shared weights (§3.5), 14 cases over `mtp_use_repeated_layer` × dense/MoE. Both MTP hazards are silent, so each needs its own guard: the async reduce-scatter path is compared numerically against the sync path on an identical model/sharding/batch, and all-gathers issued are tallied against consumes to catch a consume reading a buffer nothing gathered into. | | `test_gtp_fp8_param_gather.py` | Native-FP8 GTP_remat (§1.3): fp8-vs-BF16 loss parity (TP1/TP2, MoE), post-save-spike guard. | | `test_gtp_custom_pgs.py` | `pg_collection` plumbing: a custom `gtp_remat` group (permuted ranks, same size) must give the same fwd/bwd results as the MPU groups — catches modules reading `parallel_state` instead of the collection passed to them. | diff --git a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py index 2a2b1d63a21..2db8e62ba40 100644 --- a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py +++ b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py @@ -70,6 +70,7 @@ def model_provider(pre_process=True, post_process=True) -> GPTModel: pre_process=pre_process, post_process=post_process, fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + logit_dtype=getattr(args, 'logit_dtype', None), parallel_output=False, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, diff --git a/gpt_builders.py b/gpt_builders.py index 3512918efe6..00a9b3bd2b5 100644 --- a/gpt_builders.py +++ b/gpt_builders.py @@ -91,6 +91,7 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_ pre_process=pre_process, post_process=post_process, fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + logit_dtype=getattr(args, 'logit_dtype', None), parallel_output=True, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, diff --git a/hybrid_builders.py b/hybrid_builders.py index 7e1c58682ac..4ab668eb8bd 100644 --- a/hybrid_builders.py +++ b/hybrid_builders.py @@ -33,6 +33,7 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, pre_process=pre_process, post_process=post_process, fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + logit_dtype=getattr(args, 'logit_dtype', None), parallel_output=True, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, diff --git a/megatron/core/dist_checkpointing/strategies/fully_parallel.py b/megatron/core/dist_checkpointing/strategies/fully_parallel.py index ca7f28f3cc6..657942cadb0 100644 --- a/megatron/core/dist_checkpointing/strategies/fully_parallel.py +++ b/megatron/core/dist_checkpointing/strategies/fully_parallel.py @@ -178,6 +178,7 @@ def __init__( parallelization_group: Optional[torch.distributed.ProcessGroup] = None, do_cache_distribution: bool = False, exchange_algo: str = 'broadcast', + per_rank_object_load: bool = False, ): self.base_strategy = strategy if parallelization_group is None: @@ -187,6 +188,11 @@ def __init__( self.parallelization_group = parallelization_group self.do_cache_distribution = do_cache_distribution self.exchange_algo = exchange_algo + # When True, every rank loads *all* of its ShardedObjects directly from + # storage instead of loading only its main replicas and exchanging them + # with a WORLD-wide `all_gather_object`. Opt-in; defaults to the legacy + # gather-based exchange. See `load` for correctness rationale. + self.per_rank_object_load = per_rank_object_load self.cached_distribution: Optional[ShardDistribution] = None self.cached_global_metadata: Optional[Metadata] = None @@ -256,15 +262,47 @@ def load( assert ( len(sharded_state_dict) == 0 ), "sharded_state_dict is not empty after deferring tensors and objects" - with debug_time("base_load_ShardedObjects", logger): - # Load sharded objects first - loaded_objects = self.base_strategy.load( - to_load_objects, checkpoint_dir, async_strategy - ) - with debug_time("base_load_ShardedTensors", logger): - # Load sharded tensors separately - loaded_tensors = self.base_strategy.load(to_load_shards, checkpoint_dir, async_strategy) + if self.per_rank_object_load: + # Opt-in: every rank loads *all* of its own objects directly from + # storage, so no inter-rank object exchange is needed. This is correct + # because an object is addressed in the checkpoint by its `unique_key` + # (key + global_offset + global_shape), which excludes `replica_id`; + # each distinct object position is therefore written to disk exactly + # once (by its main replica) and any rank that needs it can read it + # locally. We merge both the main (`to_load_objects`) and non-main + # (`unloaded_objects`) replica maps so this rank loads every object in + # its state dict. This replaces the WORLD-wide `all_gather_object` + # collective with extra (but cheap) local reads of small artifacts + # (RNG states, `_extra_state`, ...). + # + # Objects are loaded together with this rank's tensor shards in a + # single base-strategy `.load()` call: `mcore_to_pyt_state_dict` + # supports a mixed tensor/object state dict, and one call means one + # metadata read and one load plan instead of two. + all_objects_to_load = {**to_load_objects, **unloaded_objects} + with debug_time("base_load_ShardedTensorsAndObjects", logger): + loaded = self.base_strategy.load( + {**to_load_shards, **all_objects_to_load}, checkpoint_dir, async_strategy + ) + # The base strategy returns the loaded values keyed by the same shard + # ids we passed in; split them back into tensors and objects. Tensor + # and object shard ids never collide (a given key is either a tensor + # or an object), so membership in the original maps is an unambiguous + # split. + loaded_tensors = {shard_id: loaded[shard_id] for shard_id in to_load_shards} + loaded_objects = {shard_id: loaded[shard_id] for shard_id in all_objects_to_load} + else: + # Default (legacy): load only this rank's main-replica objects and + # exchange them across ranks below. + with debug_time("base_load_ShardedObjects", logger): + loaded_objects = self.base_strategy.load( + to_load_objects, checkpoint_dir, async_strategy + ) + with debug_time("base_load_ShardedTensors", logger): + loaded_tensors = self.base_strategy.load( + to_load_shards, checkpoint_dir, async_strategy + ) with debug_time("self.exchange_loaded_tensors", logger): @@ -286,14 +324,17 @@ def load( with debug_time("torch.cuda.synchronize", logger): torch.cuda.synchronize() - all_loaded_objects = exchange_loaded_objects_gather_object(loaded_objects) - - if not set(unloaded_objects.keys()).issubset(all_loaded_objects.keys()): - missing_object_shards = set(unloaded_objects.keys()) - all_loaded_objects.keys() - raise CheckpointingException( - f'Missing object shards after fully parallel loading: {missing_object_shards}' - ) - torch.cuda.synchronize() + if self.per_rank_object_load: + # No object exchange: each rank already loaded every object it needs. + all_loaded_objects = loaded_objects + else: + all_loaded_objects = exchange_loaded_objects_gather_object(loaded_objects) + if not set(unloaded_objects.keys()).issubset(all_loaded_objects.keys()): + missing_object_shards = set(unloaded_objects.keys()) - all_loaded_objects.keys() + raise CheckpointingException( + f'Missing object shards after fully parallel loading: {missing_object_shards}' + ) + torch.cuda.synchronize() self.fill_in_deferred_sharded_tensors(sharded_tensors, all_loaded_tensors) self.fill_in_deferred_sharded_objects(sharded_objects, all_loaded_objects) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index a8393fcf147..284beef606e 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -564,7 +564,11 @@ def __init__( placements = Placements( dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()] ) - with fully_shard_context(device=device): + # NCCL symmetric memory requires UB. MFSDP v2 intentionally does not support UB + # without symmetric memory: it uses ncclCommRegister rather than the more performant + # ncclCommWindowRegister: + # https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html#window-registration + with fully_shard_context(device=device, use_symmetric_memory=ddp_config.nccl_ub): for submodule in reversed(list(module.modules())): if submodule is module: # The root is always sharded after selected child units so it is not @@ -642,12 +646,8 @@ def _validate_config( raise ValueError( "MFSDP v2 requires data_parallel_sharding_strategy='optim_grads_params'." ) - if ddp_config.num_distributed_optimizer_instances != 1: - raise ValueError("MFSDP v2 does not currently support HSDP.") if ddp_config.outer_dp_sharding_strategy != "no_shard": raise ValueError("MFSDP v2 does not currently support outer DP sharding.") - if ddp_config.overlap_grad_reduce or ddp_config.overlap_param_gather: - raise ValueError("MFSDP v2 does not currently support communication overlap modes.") if config.gradient_accumulation_fusion: raise ValueError("MFSDP v2 does not currently support gradient accumulation fusion.") if config.calculate_per_token_loss: @@ -657,16 +657,13 @@ def _validate_config( if config.cuda_graph_impl != "none" or ddp_config.megatron_fsdp_cuda_graph_mode: raise ValueError("MFSDP v2 does not currently support CUDA graphs.") - if ddp_config.fsdp_double_buffer: - raise ValueError("MFSDP v2 does not support fsdp_double_buffer.") if ddp_config.fsdp_db_use_persist_buf_on_alloc_fail: - raise ValueError("MFSDP v2 does not support fsdp_db_use_persist_buf_on_alloc_fail.") - if ddp_config.fsdp_all_gather_in_start_param_sync: - raise ValueError("MFSDP v2 does not support fsdp_all_gather_in_start_param_sync.") - if ddp_config.nccl_ub: - raise ValueError("MFSDP v2 does not support nccl_ub.") - if ddp_config.disable_symmetric_registration: - raise ValueError("MFSDP v2 does not support disable_symmetric_registration.") + raise ValueError( + "MFSDP v2 does not support fsdp_db_use_persist_buf_on_alloc_fail: " + "it allocates communication buffers from PyTorch memory pools." + ) + if ddp_config.nccl_ub and ddp_config.disable_symmetric_registration: + raise ValueError("MFSDP v2 requires symmetric registration when nccl_ub is enabled.") if ddp_config.fsdp_manual_registration: raise ValueError("MFSDP v2 does not support fsdp_manual_registration.") if ddp_config.delay_wgrad_compute: @@ -685,7 +682,7 @@ def _validate_config( raise ValueError("MFSDP v2 does not support megatron_fsdp_max_pool_double_buffer.") def start_param_sync(self, *unused, **unused_kwargs) -> None: - """MFSDP v2 gathers parameters from its forward pre-hook.""" + """No-op: MFSDP v2 gathers parameters from its forward pre-hooks.""" def start_grad_sync(self, *unused, **unused_kwargs) -> None: """MFSDP v2 reduces gradients during backward.""" diff --git a/megatron/core/distributed/fsdp/src/docs/mfsdp_design.md b/megatron/core/distributed/fsdp/src/docs/mfsdp_design.md new file mode 100644 index 00000000000..6bb9a93b287 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/docs/mfsdp_design.md @@ -0,0 +1,404 @@ +# Megatron FSDP Design + +Contributors: @wujingyue, @cspades, @shjwudp, @Autumn1998 + +GitHub tracker: https://github.com/orgs/NVIDIA/projects/276 + +# Executive Summary + +This design doc proposes MFSDP v2 to better satisfy +[Megatron FSDP Requirements](http://nv/mfsdp-requirements). In particular, the new +version enables: + +- **Fine-grained control.** For FSDP, automatically determining the optimal bucketing + and prefetching strategy has been challenging. + - The proposed high-level `fully_shard` API, similar to + [PyTorch FSDP2’s `fully_shard` API](https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html), + provides per-module control that is finer-grained than the production version. + - In addition, we also plan to expose lower-level APIs (e.g., ParameterGroup and + DBuffer) to give users even more fine-grained control. Spatially, this allows users + to control which parameters belong to each “bucket”. Temporally, this allows users + to control when unsharding and resharding occur during forward, backward, and + optimization. +- **Emerging optimizers**. For example, it supports tensor-atomic sharding needed for + the Muon optimizer. +- **Simplified lower-precision support.** Through a block-atomic sharding format, as + motivated by [veScale-FSDP](https://arxiv.org/abs/2602.22437). + +These capabilities are difficult to implement cleanly within the current codebase +architecture. Therefore, the code will be +[developed in branch `main`](#separate-code-paths-in-main) as a separate code path from +the existing `megatron_fsdp` implementation. The development will follow the +prototype-design-execute process that we’ll detail in +[this section](#development-process). Once the new code is on parity, we’ll gradually +migrate users over. + +The remainder of the doc focuses on the core MFSDP abstractions and building blocks that +serve as the foundation for extensions. Capabilities such as MXFP8, CUDA Graphs, double +buffering, prefetching, NCCL user buffers, HFSDP, offloading, and checkpointing build on +these primitives and introduce additional design considerations. We will cover these +areas in dedicated follow-on design documents, leveraging the interfaces and mechanisms +established here. + +# Subdesigns + +- [Optimizer](optimizer.md) +- [Runtime schedule](runtime_schedule.md) + +# API + +```py +class Placement +class Replicate(Placement) +class Partial(Placement) +class Flat(Placement) +class TensorAtomic(Placement) + + +type MeshAxis = int | str + + +@dataclass +class Placements: + dp_axes: list[MeshAxis] # outer to inner + parameter: list[Placement] # same length as dp_axes + gradient: list[Placement] + optimizer: list[Placement] + + +def fully_shard( + module: nn.Module, + mesh: DeviceMesh, + placements: Placements, + mixed_precision_policy: MixedPrecisionPolicy | None, + offload_policy: OffloadPolicy | None) -> None +``` + +Unlike MFSDP v1, `fully_shard` is expected to be called on each `nn.Module` that forms +an FSDP unit. Under the hood, `fully_shard` attaches the `FsdpModule` mixin to the +target module. FSDP units can also be nested: an outer unit owns the parameters within +its scope, excluding those managed by any inner FSDP units. This is compatible with +FSDP2’s behavior. For example, + +``` +FSDP Unit: RootModule +owns: + RootModule.root_weight + RootModule.root_bias + +contains: + FSDP Unit: SubmoduleA + owns: + SubmoduleA.weight + SubmoduleA.bias + + FSDP Unit: SubmoduleB + owns: + SubmoduleB.weight + SubmoduleB.bias + +Ownership view: + +RootModule FSDP unit +├── owns RootModule.* params +├── does NOT own SubmoduleA.* params +└── does NOT own SubmoduleB.* params + +SubmoduleA FSDP unit +└── owns SubmoduleA.* params + +SubmoduleB FSDP unit +└── owns SubmoduleB.* params +``` + +`fully_shard` sets each parameter in the given module to a shard. As a contract, no +parameters can escape its lowest FsdpModule ancestor to avoid issues like +https://github.com/NVIDIA/Megatron-LM/pull/4899. Without this contract, it can be unsafe +to unshard or reshard a parameter at the module boundary. + +The placement is similar to DTensor’s placement but for the whole **unit** and per mesh +axis. + +- `Replicate`. Not sharded. +- `Partial`. Used internally for pre-reduce-scatter gradients, which are unsharded and + only partially accumulated; not user-facing. +- `Flat`. The current per-unit, dim-0 flat sharding. Good for elementwise optimizers. +- `TensorAtomic`. Don’t cut a parameter. For emerging optimizers that need full + parameters. +- `BlockAtomic(block_size)`. Don’t cut a block of `block_size` rows. Simplifies + blockwise quantization support. Currently, a 32x1 mxfp8 block may be sharded across + ranks. This introduces complex host-side logic and custom quantization kernels to + handle two levels of absmax reduction. Using block-atomic sharding with block_size=32 + ensures that every 32-row block is owned by a single rank. +- `PerTensor(dist.tensor.Placement)`. If needed. Per-tensor dim-0 sharding used in + FSDP2. It leads to extra data copy so won’t be used by default. + +``` +Block-atomic sharding example +Input tensor: 8 rows × 4 columns +Block size: 2 rows +Earlier tensors in the parameter group occupy ranks 0, 1, and part of rank 2. + + c0 c1 c2 c3 + ┌────┬────┬────┬────┐ +r0 │ x │ x │ x │ x │ +r1 │ x │ x │ x │ x │ + ├────┼────┼────┼────┤ block 0: rows [0, 1] → rank 2 +r2 │ x │ x │ x │ x │ +r3 │ x │ x │ x │ x │ + ├────┼────┼────┼────┤ block 1: rows [2, 3] → rank 3 +r4 │ x │ x │ x │ x │ +r5 │ x │ x │ x │ x │ + ├────┼────┼────┼────┤ block 2: rows [4, 5] → rank 3 +r6 │ x │ x │ x │ x │ +r7 │ x │ x │ x │ x │ + └────┴────┴────┴────┘ block 3: rows [6, 7] → rank 4 +``` + +`Placements` encodes the various grand [sharding strategies](#sharding-strategies) we +care about and gives users the flexibility to choose which sharding format/granularity +to use for each unit. We may want to pre-define a set of common placements for +convenience, e.g., + +```py +# Assuming `Flat` placement +def hfsdp(dp_outer: MeshAxis, dp_inner: MeshAxis) -> Placements: + return Placements(dp_axes=[dp_outer, dp_inner], + parameter=[Replicate(), Flat()], + gradient=[Partial(), Flat()], + optimizer=[Flat(), Flat()]) +``` + +If needed, grouping can be customized via the `fully_shard` API, similar to +[the `buckets` argument](https://github.com/pytorch/torchtitan/pull/2378/changes#diff-35ddb8c23734307a1b5fe23e06ffe8e0f2f2c84c58943380d137371e6e21e203R3289) +in the FlexShard proposal. + +## Sharding Strategies + +Unlike MFSDP v1, MFSDP v2 does not special-case named strategies such as HSDP or HFSDP. +Instead, +`fully_shard` receives a `Placements` configuration that independently specifies the +parameter, gradient, and optimizer placement for each data-parallel axis. The table +below illustrates familiar configurations; it is not an exhaustive list of supported +strategies. + +`N` \= size of the inner DP shard dim (`dp_shard_dim`). `M` \= size of the outer DP dim +(`dp_outer_dim`), only present for HSDP/HFSDP. "Sharded" \= persistent state is +partitioned across that dim; "replicated" \= each rank holds a full copy. + +| Strategy | Parameters | Gradients | Optimizer states | +| :--------------------------------------- | :-------------------------- | :----------------------- | :-------------------------------------------------- | +| **DDP / `no_shard`** | replicated (N) | partial (N) | replicated (N) | +| **ZeRO-1 / `optim`** | replicated (N) | partial (N) | sharded (N) | +| **ZeRO-2 / `optim_grads`** | replicated (N) | sharded (N) | sharded (N) | +| **ZeRO-3 / `optim_grads_params` / FSDP** | sharded (N) | sharded (N) | sharded (N) | +| **HSDP** (FSDP inner, replicate outer) | sharded (N), replicated (M) | sharded (N), partial (M) | sharded (N), replicated (M) | +| **HFSDP** (FSDP inner, `optim` outer) | sharded (N), replicated (M) | sharded (N), partial (M) | sharded (N × M) — fully sharded across flattened DP | + +The per-axis representation also expresses combinations beyond those in the table. For +example, a configuration with ZeRO-1 on the outer DP axis and ZeRO-2 on the inner axis +uses the following placement lists, ordered to match `dp_axes` from outer to inner: + +```py +placements = Placements( + dp_axes=[dp_outer, dp_inner], + parameter=[Replicate(), Replicate()], + gradient=[Partial(), Flat()], + optimizer=[Flat(), Flat()], +) +``` + +## Compatibility + +### FSDP2 + +Introduce a separate adapter API, `fully_shard_compat`, that mirrors the signature of +PyTorch’s `fully_shard` but omits certain MFSDP-specific features. This would give +existing FSDP2 users a low-friction migration path: they can first switch to +`fully_shard_compat`, and then optionally move to MFSDP’s `fully_shard` to take +advantage of the full feature set. + +```py +def fully_shard_compat(...fsdp2 args...): + convert the args + fully_shard(...converted args...) +``` + +### MCore Adapter + +This rewrite should be mostly transparent to users of +megatron/core/distributed/fsdp/mcore_fsdp_adapter.py. We’ll implement the adapter using +the new API. + +However, certain features may behave differently. For example, +`enable_fine_grained_param_gather_hook` currently makes all-gather fine-grained (one per +submodule), but not reduce-scatter. With per-module control, users would instead apply +FSDP directly to individual submodules, causing both all-gather and reduce-scatter +operations to occur at the submodule level. + +# Key Building Blocks + +### Ownership and lifetime + +The FSDP module tree owns its persistent runtime state: + +``` +nn.Module / FsdpModule +├── active nn.Parameter +├── FsdpParameterGroup +│ ├── paired sharded and unsharded nn.Parameters +│ └── DBuffers +└── shared FsdpContext + ├── communication streams + └── prefetch-order metadata +``` + +The module’s active parameter is one of the pair owned by its parameter group. + +After construction is finalized, every backedge to the module tree **must use a weak +reference**: context prefetch metadata, parameter-group ownership markers, and hook +callbacks. Otherwise, deleting a model retains its persistent CUDA storage until cyclic +garbage collection; with weak backedges, storage is released immediately without +teardown. See https://github.com/NVIDIA/Megatron-LM/pull/6230. + +### FsdpContext + +- Created by `fully_shard_context` and shared by every `FsdpModule` constructed in that + scope. On exit, it identifies FSDP roots and finalizes the static forward and backward + prefetch orders. +- Per-device all-gather and reduce-scatter streams. Module compute runs on PyTorch’s + current stream. +- Last-microbatch state for HSDP/HFSDP gradient accumulation. +- An optional PyTorch NCCL symmetric-memory pool for communication staging buffers. + +### FsdpModule + +A mixin attached in place to the original module, so its parent retains the same child +module reference. + +- Registered forward and backward hooks drive parameter materialization, resharding, + gradient reduction, and all-gather prefetching. +- `phase` tracks the module lifecycle: `RESTING` outside module computation, `FORWARD` + between its forward hooks, and `BACKWARD` between its backward hooks. Activation + recomputation preserves `BACKWARD` through its nested forward hooks. +- Parameter groups partition the module’s owned parameters by dtype and `requires_grad`. + +### ParameterGroup + +- dtype +- requires_grad: bool +- A sharded `nn.Parameter` for every logical parameter. Its `.data` is a DTensor backed + by `main_weight`, and it is the parameter visible to the optimizer. +- The original `nn.Parameter` objects remain attached to the module. During compute, + their `.data` views a temporary replicated buffer materialized from `model_weight`; + their `.grad` is temporary full-gradient storage. +- `model_weight`: the persistent compute-dtype buffer, sharded according to + `Placements.parameter`. It may alias `main_weight` when their dtype and placements + match. +- `main_weight`: the persistent optimizer-dtype buffer, sharded according to + `Placements.optimizer`. +- `main_grad`: the persistent gradient buffer for trainable groups, sharded according to + `Placements.gradient` and allocated in the configured gradient dtype. + +### DBuffer + +Conceptually, a group of logical tensors, potentially with different shapes, stored in +one contiguous local buffer. + +- `local_buffer`: a flat `torch.Tensor` holding this rank’s contiguous shard. +- `mesh` and a per-mesh-axis `placements` tuple. The current implementation requires the + mesh to contain only data-parallel axes; callers extend returned DTensors with TP or + EP axes when needed. +- `GlobalLayout`: global tensor shapes and stable offsets used to compute every rank’s + local range. +- `redistribute(new_placements)`, with `allgather`, `allreduce`, `reduce_scatter`, and + `scatter` convenience operations. Redistributing between sharded placements preserves + the global layout; [the optimizer subdesign](optimizer.md) converts between `Flat` and + `TensorAtomic` this way. +- `get_local_tensor(index)`: the local view for one logical tensor. +- `get_dtensor(index)`: the corresponding DTensor, used by the optimizer and distributed + checkpointing. + +# Flow + +Below is what module parameters look like after each FSDP stage. + +Key contract: an FsdpModule’s owned parameters are only unsharded during its forward and +backward. + +| Stage | Action | param.data after action | param.grad after action | +| :------------------------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------ | :---------------------------- | :------------------------------------------- | +| After fully_shard / Start of a training loop | optimizer.zero_grad(set_to_none=True) | DTensor backed by main_weight | None or a zeroed DTensor backed by main_grad | +| Pre-forward / during forward | Switch to the unsharded parameter; allgather model_weight into param.data | A full-size plain Tensor | None | +| During forward | None | Unchanged | Unchanged | +| Post-forward | Release param.data; switch to the sharded parameter | DTensor backed by main_weight | None or a DTensor backed by main_grad | +| Pre-backward | Switch to the unsharded parameter; allgather model_weight into param.data | A full-size plain Tensor | None | +| During backward | Autograd sets param.grad | Unchanged | A full-size plain Tensor | +| Post-backward | Reduce-scatter param.grad; The result is written/accumulated to the sharded parameter’s grad; Release param.grad; switch to the sharded parameter | DTensor backed by main_weight | DTensor backed by main_grad | +| If more microbatches | Go back to pre-forward | Unchanged | Unchanged | +| Optimizer step | optimizer.step() | DTensor backed by main_weight | DTensor backed by main_grad | +| Post optimizer step | Quantize main_weight to model_weight | DTensor backed by main_weight | DTensor backed by main_grad | + +For HSDP and HFSDP, during post-backward for the last microbatch, we should further +reduce-scatter the sharded parameters according to `Placements.optimizer`. This way, the +reduce-scatters can be overlapped with backward compute instead of being exposed before +the optimizer step. Accordingly, prior to the forward pass of the first micro_batch, we +also need to all-gather the sharded parameters across the entire DP domain (outer \+ +inner). + +# Implementation Plan + +## Separate code paths in main + +### Production + +- [`megatron_fsdp`](../megatron_fsdp/): the existing, non-experimental implementation +- Still maintained and occasionally optimized + +### Experimental (this doc) + +- [`megatron_fsdp/experimental`](../megatron_fsdp/experimental/): the long-term version + of MFSDP that we want to maintain and use to support next generation of architectures + and training techniques +- Development will be design driven and incremental with a peer review process +- Experimental will live alongside Production **in the `main` branch** +- Once battle-tested and demonstrating performance parity (e.g. by MLPerf models) on a + per-model basis, onboard models and customers gradually. +- After enough adoption, production will become legacy and experimental will become + production + +### Prototype + +A prototype implementation by @shjwudp and @Autumn1998 remains in +[@shjwudp's fork](https://github.com/shjwudp/Megatron-LM/tree/mfsdp_refactor). The +objective is to battle-test selected features—such as per-module control and +`TracePoolAllocator`—with early users and derisk this design. Once `main` reaches the +prototype feature set (see below), further prototype development and validation will +shift to `main` so we can focus on the same code path. + +Current prototype features: + +- MXFP8 +- Overlapping +- Prefetching +- Checkpointing to DCP +- Composibility with EP +- Double buffering (through TracePoolAllocator) + +## Development process + +We’ll follow a standard prototype-design-execute process. + +1. **Prototype**: Strictly optional. Make the feature work in a draft PR **only** to + derisk the design. +2. **Design**: Update this design or create a subdesign to support a new feature. Draw + and write documentation explaining the feature and how it works. Converge and align + on the design change. +3. **Execute**: Update code and merge. Some general guidelines: + - Code, review, and test incrementally. Keep + [PRs small and focused](https://google.github.io/eng-practices/review/developer/small-cls.html). + - Favor simplicity and maintainability by default. Any performance optimization that + increases complexity should be justified with clear evidence and measurable impact. + - Critical horizontal features (for example, CUDA Graphs and `torch.compile`) should + be validated from the beginning. These integrations are easy to break and difficult + to retrofit, so we should rely on CI coverage to catch regressions early. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py index c9dc44b04db..7502cedb93a 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py @@ -41,28 +41,11 @@ ) from ..uneven_dtensor import preprocess_state_dict_for_uneven_dtensor -from .module import FsdpModule +from .parameter_group import sync_model_weights_from_main_weights __all__ = ["save_checkpoint", "load_checkpoint"] -def _sync_model_weight_from_main_weight(model: torch.nn.Module) -> None: - """Refresh every FSDP group's compute weights from its (loaded) main weights. - - A load writes into the ``main_weight``-backed sharded DTensors. When mixed precision keeps a - separate lower-precision compute buffer, that buffer is stale until the next forward pre-hook - would resync it; doing it here makes the post-load state deterministic. It is a no-op when the - compute buffer aliases the main buffer. - - Args: - model: Root module (or any module tree) containing ``FsdpModule`` instances. - """ - for module in model.modules(): - if isinstance(module, FsdpModule): - for parameter_group in module.parameter_groups: - parameter_group.sync_model_weight_from_main_weight() - - def _init_optimizer_state(optimizer: torch.optim.Optimizer) -> None: """Allocate optimizer state so a DCP load has DTensors to fill. @@ -139,4 +122,4 @@ def load_checkpoint( set_model_state_dict(model, model_state_dict) set_optimizer_state_dict(model, optimizer, optimizer_state_dict) if sync_model_weights: - _sync_model_weight_from_main_weight(model) + sync_model_weights_from_main_weights(model.parameters()) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py index 101d247ce09..4c410899285 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -148,8 +148,10 @@ class Phase(enum.Enum): # ``None`` lets pre_forward enqueue an all-gather unless an earlier FsdpModule # already prefetched this module. _unshard_event: torch.cuda.Event | None - # Backward-pre hook sets this to BACKWARD before activation recomputation - # can run. Forward and backward hooks own all other transitions. + # ``phase`` is FORWARD between pre_forward() and post_forward(), BACKWARD + # between pre_backward() and post_backward(), and RESTING otherwise. The only + # exception is non-reentrant activation recomputation: it runs between pre_backward() + # and post_backward(), preserving BACKWARD through its nested forward hooks. _phase: Phase def __init__( @@ -199,6 +201,24 @@ def context(self) -> FsdpContext: """Return the FSDP context.""" return self._context + @property + def phase(self) -> Phase: + """Return this module's lifecycle phase.""" + return self._phase + + @phase.setter + def phase(self, phase: Phase) -> None: + """Transition this module between its valid lifecycle phases.""" + allowed_transitions = { + (FsdpModule.Phase.RESTING, FsdpModule.Phase.FORWARD), + (FsdpModule.Phase.FORWARD, FsdpModule.Phase.RESTING), + (FsdpModule.Phase.RESTING, FsdpModule.Phase.BACKWARD), + (FsdpModule.Phase.BACKWARD, FsdpModule.Phase.RESTING), + } + if (self._phase, phase) not in allowed_transitions: + raise RuntimeError(f"Invalid FSDP module phase transition: {self._phase} -> {phase}.") + self._phase = phase + @property def name(self) -> str: """Return this FsdpModule's name.""" @@ -262,17 +282,14 @@ def pre_forward(self) -> None: on the comm stream, so ``AG_{i+1}`` is launched before ``F_i`` finishes. """ context = self.context + # This is the first MFSDP hook to run, so finalize the context here once + # before any module begins communication. context.ensure_finalized() - # post_forward() resets the phase after a non-recomputed forward, so a - # FORWARD phase here means this forward-pre hook ran while the previous - # forward was still in progress. - assert self._phase is not FsdpModule.Phase.FORWARD # A reentrant checkpoint recomputes before the child module's backward-pre - # hook can set its phase. Its forward still runs inside the active autograd - # GraphTask, which is the signal PyTorch FSDP2 uses as well. - is_recomputing = self._phase is FsdpModule.Phase.BACKWARD or _is_in_backward() - if not is_recomputing: - self._phase = FsdpModule.Phase.FORWARD + # hook runs. The active autograd GraphTask identifies that recomputation. + is_recomputing = self.phase is FsdpModule.Phase.BACKWARD or _is_in_backward() + if self.phase is not FsdpModule.Phase.BACKWARD: + self.phase = FsdpModule.Phase.FORWARD torch.cuda.nvtx.range_push(self._nvtx_label("forward")) self._num_ready_grad_parameters = 0 allgather_stream = context.allgather_stream @@ -317,10 +334,11 @@ def post_forward(self) -> None: # Recomputed parameters are consumed immediately by this module's # backward. Keep them materialized to avoid an unnecessary all-gather; # post_backward() will reshard them after gradient reduction. - is_recomputing = self._phase is FsdpModule.Phase.BACKWARD or _is_in_backward() + is_recomputing = self.phase is FsdpModule.Phase.BACKWARD or _is_in_backward() if not is_recomputing: self._reshard_parameter_groups() - self._phase = FsdpModule.Phase.RESTING + if self.phase is FsdpModule.Phase.FORWARD: + self.phase = FsdpModule.Phase.RESTING torch.cuda.nvtx.range_pop() def _reshard_parameter_groups(self) -> None: @@ -343,7 +361,7 @@ def _reshard_parameter_groups(self) -> None: def pre_backward(self) -> None: """Prepare full parameters and prefetch the next FsdpModule in backward order.""" - self._phase = FsdpModule.Phase.BACKWARD + self.phase = FsdpModule.Phase.BACKWARD torch.cuda.nvtx.range_push(self._nvtx_label("backward")) context = self.context current_stream = context.current_stream() @@ -370,7 +388,7 @@ def post_backward(self) -> None: """Reduce gradients and return parameters to their sharded resting state.""" self._reduce_gradient_groups() self._reshard_parameter_groups() - self._phase = FsdpModule.Phase.RESTING + self.phase = FsdpModule.Phase.RESTING torch.cuda.nvtx.range_pop() def _reduce_gradient_groups(self) -> None: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py index f1617141569..c5e4b28f211 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py @@ -14,12 +14,13 @@ """Optimizer adapter for the minimal Megatron-FSDP path.""" +from itertools import chain from typing import Any, NamedTuple import torch from torch import nn -from .parameter_group import FsdpParameterGroup, get_containing_parameter_group +from .parameter_group import get_containing_parameter_group, sync_model_weights_from_main_weights def fully_shard_optimizer( @@ -109,16 +110,9 @@ def step_post_hook( set_grad(parameter, original_grad) casted_grads.clear() - fsdp_parameter_groups: set[FsdpParameterGroup] = set() - for optimizer_group in hooked_optimizer.param_groups: - for parameter in optimizer_group["params"]: - parameter_group = get_containing_parameter_group(parameter) - if parameter_group is None: - continue - fsdp_parameter_groups.add(parameter_group) - - for parameter_group in fsdp_parameter_groups: - parameter_group.sync_model_weight_from_main_weight() + sync_model_weights_from_main_weights( + chain.from_iterable(group["params"] for group in hooked_optimizer.param_groups) + ) optimizer.register_step_pre_hook(step_pre_hook) optimizer.register_step_post_hook(step_post_hook) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 32241d8f4ea..59bd7d7b569 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -14,6 +14,7 @@ """Parameter-group runtime state for the minimal Megatron-FSDP path.""" +from collections.abc import Iterable from contextlib import nullcontext from dataclasses import dataclass from weakref import ReferenceType, ref @@ -41,6 +42,22 @@ def get_containing_parameter_group(parameter: nn.Parameter) -> "FsdpParameterGro return parameter_group_ref() +def sync_model_weights_from_main_weights(parameters: Iterable[nn.Parameter]) -> None: + """Refresh MFSDP compute weights for parameter groups represented by ``parameters``. + + Parameters outside the experimental MFSDP path are ignored. A parameter group + may own multiple parameters, but its compute-weight buffer is refreshed once. + """ + seen_parameter_groups = set() + for parameter in parameters: + if (parameter_group := get_containing_parameter_group(parameter)) is None: + continue + if parameter_group in seen_parameter_groups: + continue + seen_parameter_groups.add(parameter_group) + parameter_group.sync_model_weight_from_main_weight() + + @dataclass(frozen=True, eq=False) class FsdpParameter: """One physical parameter and its FSDP runtime representations.""" diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index e806a9fdd09..1f67b1a5b5b 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1749,7 +1749,7 @@ class TELMHeadColumnParallelLinear(TEColumnParallelLinear): ``delay_wgrad_compute`` is forced off to mirror its no-op ``backward_dw``, and ``get/set_extra_state`` match the bf16 LM head's state-dict shim. The LM-head kwargs ``keep_master_weight_for_test``, ``skip_weight_param_allocation``, - ``defer_embedding_wgrad_compute`` buffers, and ``disable_grad_reduce`` are + ``defer_embedding_wgrad_compute`` buffers, ``disable_grad_reduce``, and ``output_dtype`` are accepted to preserve the ``ColumnParallelLinear`` signature but currently raise when set non-default — TE will not support them natively, so they would have to be implemented in this subclass, which has not been done yet. @@ -1776,6 +1776,7 @@ def __init__( tp_comm_buffer_name: Optional[str] = None, disable_grad_reduce: bool = False, tp_group: Optional[torch.distributed.ProcessGroup] = None, + output_dtype: Optional[torch.dtype] = None, ): from megatron.core.fp8_utils import is_mxfp8_output_proj_active @@ -1794,6 +1795,8 @@ def __init__( ) if disable_grad_reduce: raise ValueError("TE output projection does not support disable_grad_reduce.") + if output_dtype is not None: + raise ValueError("TE MXFP8 output projection does not support output_dtype.") te_config = copy.copy(config) # Match ColumnParallelLinear.backward_dw's no-op so the LM head keeps diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py index 35fa97b3d38..bda35c3993a 100644 --- a/megatron/core/models/common/model_chunk_schedule_plan.py +++ b/megatron/core/models/common/model_chunk_schedule_plan.py @@ -351,7 +351,7 @@ def __init__( loss_mask: Optional[Tensor] = None, padding_mask=None, *, - output_processor: Optional[Callable[..., Tensor]] = None, + output_processor: Optional[Callable[..., Any]] = None, output_processor_context: Optional[Any] = None, ): """Initialize the schedule plan of all Transformer layers' sub-modules. diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index ff3514b7433..ad5f46c1f97 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -66,6 +66,9 @@ class GPTModel(LanguageModule): Include an output layer (used with pipeline parallelism). Defaults to True. fp16_lm_cross_entropy (bool, optional): Defaults to False. + logit_dtype (torch.dtype, optional): + Dtype for the output-layer GEMM result. Defaults to None, which uses + the hidden-state dtype. parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor parallel ranks. Defaults to True. @@ -100,6 +103,7 @@ def __init__( pre_process: bool = True, post_process: bool = True, fp16_lm_cross_entropy: bool = False, + logit_dtype: Optional[torch.dtype] = None, parallel_output: bool = True, share_embeddings_and_output_weights: bool = False, position_embedding_type: Literal[ @@ -133,6 +137,7 @@ def __init__( self.pre_process = pre_process self.post_process = post_process self.fp16_lm_cross_entropy = fp16_lm_cross_entropy + self.logit_dtype = logit_dtype self.parallel_output = parallel_output self.share_embeddings_and_output_weights = share_embeddings_and_output_weights self.vp_stage = vp_stage @@ -283,6 +288,7 @@ def __init__( embedding_activation_buffer=self.embedding_activation_buffer, grad_output_buffer=self.grad_output_buffer, tp_group=self.pg_collection.tp, + output_dtype=self.logit_dtype, ) if self.pre_process or self.post_process or self.mtp_process: @@ -531,9 +537,9 @@ def forward( inference_params: Optional[BaseInferenceContext] = None, loss_mask: Optional[Tensor] = None, padding_mask: Optional[Tensor] = None, - output_processor: Optional[Callable[..., Tensor]] = None, + output_processor: Optional[Callable[..., Any]] = None, output_processor_context: Optional[Any] = None, - ) -> Tensor: + ) -> Any: """Forward function of the GPT Model This function passes the input tensors through the embedding layer, and then the decoder and finally into the post processing layer (optional). @@ -810,7 +816,7 @@ def build_schedule_plan( loss_mask: Optional[Tensor] = None, padding_mask: Optional[Tensor] = None, *, - output_processor: Optional[Callable[..., Tensor]] = None, + output_processor: Optional[Callable[..., Any]] = None, output_processor_context: Optional[Any] = None, ): """Builds a computation schedule plan for the model. diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index f0358de57b9..65ab3ba1178 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -3,6 +3,7 @@ import logging from typing import Literal, Optional +import torch from torch import Tensor from megatron.core import tensor_parallel @@ -83,6 +84,8 @@ class HybridModel(LanguageModule, GraphableMegatronModule): post_process (bool, optional): Include an output layer (used with pipeline parallelism). Defaults to True. fp16_lm_cross_entropy (bool, optional): Defaults to False. + logit_dtype (torch.dtype, optional): Dtype for the output-layer GEMM result. + Defaults to None, which uses the hidden-state dtype. parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor parallel ranks. Defaults to True. share_embeddings_and_output_weights (bool, optional): When True, input embeddings and @@ -113,6 +116,7 @@ def __init__( pre_process: bool = True, post_process: bool = True, fp16_lm_cross_entropy: bool = False, + logit_dtype: Optional[torch.dtype] = None, parallel_output: bool = True, share_embeddings_and_output_weights: bool = False, # Mamba with no attention has no need for position embeddings, so none is default @@ -144,6 +148,7 @@ def __init__( self.pre_process = pre_process self.post_process = post_process self.fp16_lm_cross_entropy = fp16_lm_cross_entropy + self.logit_dtype = logit_dtype self.parallel_output = parallel_output self.share_embeddings_and_output_weights = share_embeddings_and_output_weights self.position_embedding_type = position_embedding_type @@ -323,6 +328,7 @@ def __init__( skip_weight_param_allocation=self.pre_process and self.share_embeddings_and_output_weights, tp_group=self.pg_collection.tp, + output_dtype=self.logit_dtype, pg_collection=self.pg_collection, ) diff --git a/megatron/core/optimizer/fully_sharded_optimizer.py b/megatron/core/optimizer/fully_sharded_optimizer.py index 18c2354dcb2..9c70c5b6f86 100644 --- a/megatron/core/optimizer/fully_sharded_optimizer.py +++ b/megatron/core/optimizer/fully_sharded_optimizer.py @@ -8,6 +8,9 @@ from ..config_logger import has_config_logger_enabled, log_config_to_disk from ..dist_checkpointing.mapping import ShardedStateDict +from ..distributed.fsdp.src.megatron_fsdp.experimental.parameter_group import ( + sync_model_weights_from_main_weights, +) from ..transformer.module import MegatronModule from .grad_scaler import MegatronGradScaler from .optimizer import MixedPrecisionOptimizer @@ -120,7 +123,8 @@ def _copy_model_grads_to_main_grads(self) -> None: """No-op: MFSDP v2 reduces directly into optimizer-visible sharded grads.""" def _copy_main_params_to_model_params(self) -> None: - """No-op: MFSDP v2 currently syncs compute weights in its forward pre-hook.""" + """Refresh MFSDP V2 compute weights after updating optimizer weights.""" + sync_model_weights_from_main_weights(self.get_parameters()) def _copy_model_params_to_main_params(self, state_dict=None) -> None: """No-op: model loads already write into MFSDP v2's main weights.""" diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index d1d792ca76d..a791d262954 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -213,15 +213,19 @@ def _compute_per_buffer_param_layout( followed by an isolated bucket for that embedding alone. 2. When the chunk's total numel reaches ``bucket_size`` (or all params have been consumed), bin-pack the chunk into ``dp_size`` - shards via greedy LPT — sort by numel descending and assign each - param to the shard with the smallest current load. + shards: sort by estimated Newton-Schulz compute cost descending and + assign each param to the shard with the smallest accumulated compute + load, subject to a per-bucket numel cap that bounds shard-imbalance + padding. Compute loads persist across buckets so expensive + (GTP-sharded) matrices spread over the whole buffer instead of + clustering inside each bucket. 3. Pad each shard to ``max(shard_cursors)`` aligned to :meth:`_shard_divisor`, then emit the bucket. Each bucket therefore spans a contiguous backprop range so that ``overlap_grad_reduce`` can dispatch the bucket's reduce-scatter as soon as the bucket's backward segment finishes — preserving the - original DDP overlap semantics. LPT bin-packing keeps shards close + original DDP overlap semantics. Greedy bin-packing keeps shards close to balanced; for uniform transformer blocks where ``params_per_layer * num_layers`` is a multiple of ``dp_size`` the packing is perfect. @@ -247,6 +251,27 @@ def _compute_per_buffer_param_layout( bucket_id = 0 shard_imbalance_padding_numel = 0 + # Persistent compute loads across buckets so LPT spreads expensive + # (GTP-sharded) params evenly instead of clustering them per bucket. + shard_compute_loads = [0] * dp_size + + def _ns_compute_cost(param): + """Estimate Newton-Schulz compute cost for a parameter. + + Newton-Schulz only runs on matrices, so anything that is not 2D falls + back to its element count. For a 2D param the cost is + ~ max(M,N) * min(M,N)^2, the dominant term in the orthogonalization; + GTP-sharded params reconstruct the full post-AllGather shape first + (GTP always shards along dim 0). + """ + if param.dim() != 2: + return param.data.nelement() + m, n = param.data.shape + if getattr(param, 'is_gtp_weight_remat', False): + m = m * getattr(param, 'gtp_remat_size', 1) + big, small = max(m, n), min(m, n) + return big * small * small + def _emit_bucket( chunk_params: List[torch.nn.Parameter], shared_embedding: bool = False ) -> None: @@ -276,17 +301,33 @@ def _emit_bucket( shard_assignments[shard_id].append((None, numel)) shard_cursors[shard_id] = numel else: - # Greedy LPT: largest first, assign to the least-loaded shard. - # The within-shard order is sorted-by-numel, not backprop — + # Compute-balanced LPT: sort by Newton-Schulz compute cost + # (accounts for full post-AllGather shape under GTP), assign to + # the shard with least accumulated compute load. Compute loads + # persist across buckets; numel cursors reset per bucket. + # A per-bucket numel cap prevents excessive padding. + # The within-shard order is sorted-by-compute-cost, not backprop; # that is fine because all params in the chunk share the same # bucket_id, so DDP's backprop-order iteration still sees # monotonic bucket_ids across the chunk boundary. - for param in sorted(chunk_params, key=lambda p: -p.data.nelement()): + _NUMEL_EPSILON = 0.3 + total_chunk_numel = sum(p.data.nelement() for p in chunk_params) + max_shard_numel = total_chunk_numel / dp_size * (1 + _NUMEL_EPSILON) + for param in sorted(chunk_params, key=lambda p: -_ns_compute_cost(p)): numel = param.data.nelement() - min_shard = min(range(dp_size), key=lambda s: shard_cursors[s]) + candidates = [ + s + for s in range(dp_size) + if pad_param_start(shard_cursors[s]) + numel <= max_shard_numel + ] + if candidates: + min_shard = min(candidates, key=lambda s: shard_compute_loads[s]) + else: + min_shard = min(range(dp_size), key=lambda s: shard_cursors[s]) placement = pad_param_start(shard_cursors[min_shard]) shard_assignments[min_shard].append((param, numel)) shard_cursors[min_shard] = placement + numel + shard_compute_loads[min_shard] += _ns_compute_cost(param) padded_shard_size = pad_to_divisor(max(shard_cursors), shard_divisor) bucket_start_index = buffer_cursor @@ -325,8 +366,18 @@ def _emit_bucket( chunk_params: List[torch.nn.Parameter] = [] chunk_numel = 0 chunk_max_param = 0 - # Mirror _emit_bucket's greedy LPT placement incrementally so we can - # decide, per param, whether it still fits in the current bucket. + # Approximate _emit_bucket's placement so we can decide, per param, + # whether it still fits in the current bucket. This estimates rather + # than mirrors, for two reasons: _emit_bucket sorts the chunk before + # packing it, while this places params in backprop order, and + # _emit_bucket assigns by Newton-Schulz compute cost, while this tracks + # numel. Equal-sized params make both differences vanish, because + # sorting is then a no-op and cost is proportional to numel. Mixed + # sizes send params to different shards under the two orders, so the + # real maximum shard load can exceed the estimated one. _absorbs then + # admits a param that does grow the bucket, leaving a buffer larger + # than closing the bucket early would have produced. The layout stays + # valid either way; see test_mixed_sizes_can_absorb_into_larger_bucket. shard_loads = [0] * dp_size def _absorbs(numel: int) -> bool: diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 1baaea0ea3a..3fa69813680 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -21,6 +21,9 @@ ) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gdp_context_parallel import GDPContextParallel + +# Decode uses the in-repo Triton conv update, which accepts int64 slot indices. +from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update from megatron.core.ssm.packed_seq_helpers import ( build_packed_seq_idx, check_fla_sequence_packing_support, @@ -38,11 +41,10 @@ from megatron.core.utils import deprecate_inference_params try: - from causal_conv1d import causal_conv1d_fn, causal_conv1d_update + from causal_conv1d import causal_conv1d_fn from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states except ImportError: causal_conv1d_fn = None - causal_conv1d_update = None causal_conv1d_varlen_states = None try: @@ -604,9 +606,8 @@ def ssm_decode( dim=-1, ) - # Indexed conv update: reads/writes the per-request conv state rows - # selected by ``batch_indices``, in place. ``self.activation`` must be the - # activation *string* so the kernel enables SiLU (a bool would disable it). + # Indexed conv update into the per-request state rows (``batch_indices`` + # is None for static batching, where the cache is already in order). VKQ = causal_conv1d_update( VKQ, conv_state, diff --git a/megatron/core/ssm/gdn_layer_config.py b/megatron/core/ssm/gdn_layer_config.py new file mode 100644 index 00000000000..15f509250f3 --- /dev/null +++ b/megatron/core/ssm/gdn_layer_config.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.transformer.transformer_config import TransformerConfig + + +class GDNLayerConfig(TransformerConfig): + """Configuration for a Gated DeltaNet layer in a hybrid stack. + + Due to backwards-compatibility, this config's arguments are defined in TransformerConfig. + """ diff --git a/megatron/core/ssm/mamba_layer_config.py b/megatron/core/ssm/mamba_layer_config.py new file mode 100644 index 00000000000..a9b58c3843d --- /dev/null +++ b/megatron/core/ssm/mamba_layer_config.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.transformer.transformer_config import TransformerConfig + + +class MambaLayerConfig(TransformerConfig): + """Configuration for a Mamba layer in a hybrid stack. + + Due to backwards-compatibility, this config's arguments are defined in TransformerConfig. + """ diff --git a/megatron/core/ssm/mlp_layer_config.py b/megatron/core/ssm/mlp_layer_config.py new file mode 100644 index 00000000000..6da7de589d2 --- /dev/null +++ b/megatron/core/ssm/mlp_layer_config.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.transformer.transformer_config import TransformerConfig + + +class MLPLayerConfig(TransformerConfig): + """Configuration for a dense MLP layer in a hybrid stack. + + Due to backwards-compatibility, this config's arguments are defined in TransformerConfig. + """ diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index 5eac911d1b3..b8d2a8faff2 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -217,6 +217,10 @@ def _classify_param_chain(param_name: str) -> str: if not scope: # CG disabled return U + # MoE latent projections. + if ".mlp.fc1_latent_proj." in n or ".mlp.fc2_latent_proj." in n: + return G if "moe_router" in scope else U + if ".mlp.shared_experts." in n: if _MOE_SHARED_EXPERT_OVERLAP: return U @@ -878,6 +882,11 @@ def _init_gtp_runtime_attrs(obj): obj._recompute_prefetch_handle = None obj._recompute_ag_event = torch.cuda.Event(external=True) obj._recompute_already_drained = False + # Own AG buffer for the recompute chain, with its own parity so a one-ahead prefetch cannot + # land in the buffer the previous recompute node is still reading. See + # _ensure_no_shared_buffer_with. + obj._ag_ticket_recompute = None + obj._recompute_buf_parity = None # Chain identity (GRAPHED/UNGRAPHED). Defaults to UNGRAPHED; classify_gtp_chains(model) # walks the model at init (after set_cuda_graph_modules) and reclassifies on param name + # active cuda_graph_modules. @@ -926,6 +935,7 @@ class GTPShardedParam(torch.nn.Parameter): _recompute_chain_state: Dict[str, dict] = {} _link_tables_flushed: bool = False + _recompute_link_tables_flushed: bool = False @classmethod def _get_chain_state(cls, chain_id: str) -> dict: @@ -940,7 +950,11 @@ def _get_chain_state(cls, chain_id: str) -> dict: @classmethod def _get_recompute_chain_state(cls, chain_id: str) -> dict: if chain_id not in cls._recompute_chain_state: - cls._recompute_chain_state[chain_id] = {"last_weight": None} + cls._recompute_chain_state[chain_id] = { + "last_weight": None, + "link_node_count": 0, + "link_table_buffer": [], + } return cls._recompute_chain_state[chain_id] @classmethod @@ -950,18 +964,50 @@ def flush_link_tables(cls) -> None: Call only where the chains are complete -- NOT on "this weight is already linked", which MTP hits mid-forward while later links are still being created. """ + # Clear each buffer on emit: the latch alone is not enough, since the dynamic GTP_* + # subclasses each carry their own copy of it over this one shared buffer set. + emitted = False for chain in cls._chain_state.values(): if chain["link_table_buffer"]: log_single_rank(logger, logging.INFO, "\n".join(chain["link_table_buffer"]) + "\n") - cls._link_tables_flushed = True + chain["link_table_buffer"] = [] + emitted = True + cls._link_tables_flushed = emitted + + @classmethod + def flush_recompute_link_tables(cls) -> None: + """Log every recompute chain's link table once, atomically. + + Recompute links are built during backward, so call this only where one has finished -- + never on "this weight is already linked", which a replayed MTP block reaches too early. + """ + # Latch on having emitted, not on having been called: the first forward runs before any + # backward has built a chain, and must not suppress the real flush. + emitted = False + for rchain in cls._recompute_chain_state.values(): + if rchain["link_table_buffer"]: + log_single_rank(logger, logging.INFO, "\n".join(rchain["link_table_buffer"]) + "\n") + rchain["link_table_buffer"] = [] + emitted = True + cls._recompute_link_tables_flushed = emitted + + @classmethod + def _recompute_link_table_row( + cls, prev: "GTPShardedParam", curr: "GTPShardedParam", rchain: dict + ) -> None: + """Buffer one recompute-chain link row, under its own table heading.""" + cls._buffer_link_table_row(prev, curr, rchain, label="RECOMPUTE chain") @classmethod def _buffer_link_table_row( - cls, prev: "GTPShardedParam", curr: "GTPShardedParam", chain: dict + cls, prev: "GTPShardedParam", curr: "GTPShardedParam", chain: dict, label: str = "chain" ) -> None: - """Buffer one prefetch-link row (flushed atomically on the second forward pass).""" + """Buffer one prefetch-link row (flushed atomically once the chain is complete). + + ``label`` lets the recompute chain reuse this with its own table heading. + """ _W = 70 - _D = 20 + _D = 8 # widest realistic value is "bfloat16"; MXFP8/NVFP4 are 5 _S = 20 def _layer_id(name: str) -> str: @@ -977,40 +1023,37 @@ def _shape(param: "GTPShardedParam") -> str: return str(tuple(param.shape)) def _dtype(param: "GTPShardedParam") -> str: - # Report the dtype of the tensor that is ACTUALLY all-gathered, not the - # GTPShardedParam wrapper (whose logical dtype is the high-precision model-weight - # shard, i.e. params_dtype — bf16 in mixed precision). When the param has an FP8 - # representation (``param.quantized`` populated — by --fp8-param-gather's optimizer - # FP32->FP8 write, or by the per-forward cast otherwise), that quantized tensor is - # what gets gathered, yet a TE QuantizedTensor still reports a "fake" params_dtype - # ``.dtype``. So surface its raw storage dtype (e.g. uint8) tagged with the quantized - # class to make the FP8 all-gather unambiguous. + # ``.dtype`` lies here: the wrapper and the TE quantized tensor both report + # params_dtype (bf16), so read the actually-gathered format off the quantized class. q = getattr(param, "quantized", None) if getattr(param, "_gtp_native_fp8", False) and q is not None: - raw = getattr(q, "_rowwise_data", None) - if raw is None: - raw = getattr(q, "_data", None) - raw_dt = str(raw.dtype).replace("torch.", "") if raw is not None else "?" - return f"{type(q).__name__}/{raw_dt}" - return str(getattr(param, "dtype", "-")) + # GTP_MXFP8Tensor -> MXFP8. Derived, not hardcoded, so NVFP4/FP8 recipes work too. + name = type(q).__name__ + if name.startswith("GTP_"): + name = name[len("GTP_") :] + for suffix in ("QTensor", "Tensor"): + if name.endswith(suffix): + name = name[: -len(suffix)] + break + return name + return str(getattr(param, "dtype", "-")).replace("torch.", "") chain["link_node_count"] += 1 if chain["link_node_count"] == 1: chain_id = getattr(curr, "chain_id", GTPChain.UNGRAPHED.value) chain["link_table_buffer"].append( - f"\n[{chain_id} chain]\n{'node_id':>7} | {'layer_id':>8} |" - f" {'dtype':<{_D}} | {'shape':<{_S}} | {'curr_weight_name':<{_W}} |" - f" prev_weight_name\n{'-'*7}-+-{'-'*8}-+-{'-'*_D}-+-{'-'*_S}-+-{'-'*_W}-+-{'-'*_W}" + f"\n[{chain_id} {label}]\n{'node_id':>7} | {'layer_id':>8} |" + f" {'dtype':<{_D}} | {'shape':<{_S}} | weight_name\n" + f"{'-'*7}-+-{'-'*8}-+-{'-'*_D}-+-{'-'*_S}-+-{'-'*_W}" ) - # Seed weight (first GTP param) as row 0 + # Seed weight (chain head) as row 0 chain["link_table_buffer"].append( f"{'0':>7} | {_layer_id(prev._debug_name):>8} | " - f"{_dtype(prev):<{_D}} | {_shape(prev):<{_S}} | {prev._debug_name:<{_W}} | -" + f"{_dtype(prev):<{_D}} | {_shape(prev):<{_S}} | {prev._debug_name}" ) chain["link_table_buffer"].append( f"{chain['link_node_count']:>7} | {_layer_id(curr._debug_name):>8} | " - f"{_dtype(curr):<{_D}} | {_shape(curr):<{_S}} | " - f"{curr._debug_name:<{_W}} | {prev._debug_name}" + f"{_dtype(curr):<{_D}} | {_shape(curr):<{_S}} | {curr._debug_name}" ) @staticmethod @@ -1089,30 +1132,35 @@ def _gather_buffer_identity(self, dtype) -> tuple: """The part of the cache key that decides which weights share a gather buffer.""" return (self._unsharded_shape_padded, dtype, self.expert_idx) - def _ensure_distinct_buffer_from_prev(self, dtype): - """Move self to a second buffer if its chain predecessor would share one. + def _ensure_no_shared_buffer_with(self, predecessor, predecessor_dtype, dtype, parity_attr): + """Guarantee that two adjacent weights on a prefetch chain never share a gather buffer. + + Sharing one is a data race: one-step-ahead prefetch keeps both neighbours live at once, + so self's gather writes the buffer while the predecessor's GEMM is still reading it. They + share a buffer exactly when they resolve to the same cache key (same gathered shape and + dtype); flipping ``parity_attr`` moves self to a second buffer and breaks the tie. + Differently-shaped neighbours never shared a key, so this is a no-op for them. - One-step-ahead prefetch keeps prev_w and self live at once, so sharing a buffer lets - self's gather clobber the weight prev_w's GEMM is still reading. Neighbours normally - differ in shape; a CUDA-graph-partitioned chain can leave two same-shaped weights - adjacent (embedding + output_layer alone in the UNGRAPHED chain). + Which chain to guard is the caller's to say: pass (``prev_w``, ``_buf_parity``) or + (``_recompute_prev``, ``_recompute_buf_parity``). No default, because the chains disagree + on who a weight's neighbour is. ``predecessor_dtype`` is the dtype that weight actually + gathers in, ``None`` if it never has -- passed in rather than read off the predecessor, + because grouped weights cache their dtypes on the batch anchor, not per expert. - Grouped chains use their own counter (``_GTP_GROUPED_BUF_PARITY_COUNTER``). + Callers on the fwd chain skip grouped weights, which get their parity from + ``_GTP_GROUPED_BUF_PARITY_COUNTER`` instead. """ - prev = self.prev_w - if prev is None or _chain_is_grouped(self.chain_id): - return - if self.is_routed_expert or prev.is_routed_expert: - return - if prev._cached_dtypes is None: # never gathered — no buffer to collide with + if predecessor is None or predecessor_dtype is None: # nothing gathered to collide with return - if prev._gather_buffer_identity(prev._cached_dtypes[0]) != self._gather_buffer_identity( + if predecessor._gather_buffer_identity(predecessor_dtype) != self._gather_buffer_identity( dtype ): return - self._buf_parity = 1 - (getattr(prev, "_buf_parity", None) or 0) + setattr(self, parity_attr, 1 - (getattr(predecessor, parity_attr, None) or 0)) - def _get_cache_key(self, dtype, fwd: bool, reduce_scatter: bool) -> tuple: + def _get_cache_key( + self, dtype, fwd: bool, reduce_scatter: bool, recompute: bool = False + ) -> tuple: """Build a cache key that includes the communication scheduling domain. ``GTPWeightCache.release`` retains a ticket's buffer pointer while returning the storage to @@ -1155,9 +1203,16 @@ def _get_cache_key(self, dtype, fwd: bool, reduce_scatter: bool) -> tuple: # parity alternates consecutive blocks between two buffers. key = key + (self.chain_id, self._double_buffer_parity()) elif getattr(self, "_buf_parity", None): - # Set by _ensure_distinct_buffer_from_prev. Parity 0 keeps the shared buffer, so + # Set by _ensure_no_shared_buffer_with. Parity 0 keeps the shared buffer, so # only the second weight of an adjacent same-key pair costs an extra allocation. key = key + (self._buf_parity,) + if recompute: + # Two components, guarding two different collisions: + # "recompute" keeps these buffers away from the fwd ones, which may still hold a + # prefetch in flight when a recompute gather lands; + # the parity keeps recompute NEIGHBOURS apart. _buf_parity above cannot do that -- + # it is decided against prev_w, and the recompute chain links different weights. + key = key + ("recompute", getattr(self, "_recompute_buf_parity", None) or 0) return key def _strip_padding(self, tensor): @@ -1213,8 +1268,15 @@ def _strip_padding(self, tensor): return tensor[: -self.pad_length] - def _all_gather_weight(self, async_op, fwd, nvtx_label=None): - """Quantize (if needed) and all-gather weight. Returns (weight_total, handle).""" + def _all_gather_weight( + self, async_op, fwd, nvtx_label=None, recompute=False, recompute_prev=None + ): + """Quantize (if needed) and all-gather weight. Returns (weight_total, handle). + + ``recompute=True`` targets the recompute chain's own buffer, not the fwd/bwd one; + ``recompute_prev`` is this node's recompute-chain predecessor, used once to pick a + non-colliding buffer. + """ if nvtx_label is None: nvtx_label = ( self._debug_name + (".fwd" if fwd else ".bwd") + (".async" if async_op else ".sync") @@ -1260,8 +1322,27 @@ def _all_gather_weight(self, async_op, fwd, nvtx_label=None): self._cached_dtypes = dtypes out_buffers = [] cache = get_global_GTP_cache() - for p, dt in zip(weights, dtypes): - if fwd: + # Match experts index-for-index: the cache key carries expert_idx, so expert k collides + # with expert k of the neighbouring block, never with that block's anchor. + prev_weights = recompute_prev._weights if recompute_prev is not None else [] + prev_dtypes = recompute_prev._cached_dtypes if recompute_prev is not None else None + for idx, (p, dt) in enumerate(zip(weights, dtypes)): + if recompute: + if p._ag_ticket_recompute is None: + # Must run before reserve — it decides which buffer the ticket gets. + p._ensure_no_shared_buffer_with( + predecessor=prev_weights[idx] if idx < len(prev_weights) else None, + predecessor_dtype=( + prev_dtypes[idx] if prev_dtypes and idx < len(prev_dtypes) else None + ), + dtype=dt, + parity_attr="_recompute_buf_parity", + ) + p._ag_ticket_recompute = cache.reserve(p, dt, fwd=True, recompute=True) + cache.get(p._ag_ticket_recompute) + cache.release(p._ag_ticket_recompute) + out_buffers.append(cache.get(p._ag_ticket_recompute)) + elif fwd: if p._ag_ticket_fwd is None: p._ag_ticket_fwd = cache.reserve(p, dt, fwd=True) cache.get(p._ag_ticket_fwd) @@ -1348,8 +1429,15 @@ def _wait_param_gather(self): self._prefetch_handle = None self.ag_event.record() - def _all_gather_weight_on_demand(self, fwd): - result, _ = self._all_gather_weight(async_op=False, fwd=fwd) + def _all_gather_weight_on_demand(self, fwd, recompute=False, recompute_prev=None): + # Only pass the recompute kwargs when they apply, so the fwd/bwd path keeps calling + # _all_gather_weight with its original signature. + if recompute: + result, _ = self._all_gather_weight( + async_op=False, fwd=fwd, recompute=True, recompute_prev=recompute_prev + ) + else: + result, _ = self._all_gather_weight(async_op=False, fwd=fwd) result = result if self.is_routed_expert else [result] result = [self._strip_padding(r) for r in result] result = [r.detach().requires_grad_(w.requires_grad) for r, w in zip(result, self._weights)] @@ -1420,14 +1508,16 @@ def _wait_recompute_param_gather(self): def _recompute_prefetch_next(self, target, nvtx_label=None): # Issue target's rowwise (fwd) AG into its recompute slot. _all_gather_weight skips the - # AG-state transition under recompute, so target's dgrad state is untouched; result lands - # in target._ag_ticket_fwd. - _, handle = target._all_gather_weight(async_op=True, fwd=True, nvtx_label=nvtx_label) + # AG-state transition under recompute, so target's dgrad state is untouched; the write + # lands in target._ag_ticket_recompute, which self is guaranteed not to be reading. + _, handle = target._all_gather_weight( + async_op=True, fwd=True, nvtx_label=nvtx_label, recompute=True, recompute_prev=self + ) target._recompute_prefetch_handle = handle def _get_recompute_prefetched_weight(self): # Recompute-chain analogue of _get_prefetched_weight (state-neutral; reads the - # rowwise _ag_ticket_fwd via the _recompute_* slot). + # rowwise gather via the _recompute_* slot). if self._recompute_already_drained: # Producer already drained via wait_async_comms (CG capture); skip the # captured cross-graph wait (CUDA no-op anyway). @@ -1439,7 +1529,7 @@ def _get_recompute_prefetched_weight(self): result = [] cache = get_global_GTP_cache() for w in self._weights: - result.append(cache.get(w._ag_ticket_fwd)) + result.append(cache.get(w._ag_ticket_recompute)) result = [self._strip_padding(r) for r in result] result = [r.detach().requires_grad_(w.requires_grad) for r, w in zip(result, self._weights)] return result if self.is_routed_expert else result[0] @@ -1508,6 +1598,11 @@ def all_gather_and_prefetch(self, fwd: bool = True, nvtx_label: str = None): in_recompute = in_fp8_activation_recompute_phase() use_recompute_chain = in_recompute and GTP_CONFIG.weight_prefetch + # Reaching a forward gather proves the previous backward finished, so the chains it + # built are complete. Mirrors flush_link_tables, which fires on the first backward AG. + if not in_recompute and not type(self)._recompute_link_tables_flushed: + type(self).flush_recompute_link_tables() + # Consume current weight. if use_recompute_chain and self._recompute_prev is not None: result = self._get_recompute_prefetched_weight() @@ -1518,8 +1613,20 @@ def all_gather_and_prefetch(self, fwd: bool = True, nvtx_label: str = None): and self._prefetch_available() ): result = self._get_prefetched_weight(True) + elif use_recompute_chain: + # Recompute chain head. It still needs the recompute buffer, and on the first + # backward the chain links do not exist yet, so take the predecessor from the cursor. + result = self._all_gather_weight_on_demand( + True, + recompute=True, + recompute_prev=( + self._recompute_prev + or type(self)._get_recompute_chain_state(self.chain_id)["last_weight"] + ), + ) else: - # On-demand: chain head (fwd or recompute global-first) or first-iter build. + # On-demand: fwd chain head or first-iter build. Deliberately called with the + # original signature so the recompute plumbing never perturbs the fwd path. result = self._all_gather_weight_on_demand(True) # Prefetch next weight on the matching chain. @@ -1559,6 +1666,8 @@ def all_gather_and_prefetch(self, fwd: bool = True, nvtx_label: str = None): if last_r is not None and last_r._recompute_next is None: last_r._recompute_next = self self._recompute_prev = last_r + # Only once a link exists, so the head lands in row 0 -- same as the fwd table. + cls._recompute_link_table_row(last_r, self, rchain) self._recompute_initialized = True rchain["last_weight"] = self @@ -1581,7 +1690,19 @@ def all_gather_and_prefetch(self, fwd: bool = True, nvtx_label: str = None): q.dtype if q is not None else w.dtype for q, w in zip(quantizers, self._weights) ] # Must run before the reserve below — it decides which buffer the ticket gets. - self._ensure_distinct_buffer_from_prev(dtypes[0]) + # Grouped/routed weights take their fwd parity from _double_buffer_parity instead. + prev_w = self.prev_w + if not _chain_is_grouped(self.chain_id) and not self.is_routed_expert: + self._ensure_no_shared_buffer_with( + predecessor=prev_w, + predecessor_dtype=( + prev_w._cached_dtypes[0] + if prev_w is not None and prev_w._cached_dtypes + else None + ), + dtype=dtypes[0], + parity_attr="_buf_parity", + ) for w, dt in zip(self._weights, dtypes): w._ag_ticket_fwd = cache.reserve(w, dt, fwd=True) @@ -2121,9 +2242,11 @@ def _allocate_buffer( ) return buf - def reserve(self, param: "GTPShardedParam", dtype, fwd: bool, reduce_scatter=False) -> int: + def reserve( + self, param: "GTPShardedParam", dtype, fwd: bool, reduce_scatter=False, recompute=False + ) -> int: """Assign a persistent ticket. No buffer is allocated until ``get()``.""" - key = param._get_cache_key(dtype, fwd, reduce_scatter) + key = param._get_cache_key(dtype, fwd, reduce_scatter, recompute=recompute) ticket = self._next_ticket self._next_ticket += 1 @@ -2373,6 +2496,7 @@ def reset_gtp_state(): GTPShardedParam._chain_state.clear() GTPShardedParam._recompute_chain_state.clear() GTPShardedParam._link_tables_flushed = False + GTPShardedParam._recompute_link_tables_flushed = False _GTP_GROUPED_BUF_PARITY_COUNTER.clear() diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 0f248bcf399..e842e5d9f39 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -411,21 +411,20 @@ class LinearWithFrozenWeight(torch.autograd.Function): @staticmethod @custom_fwd - def forward(ctx, input, weight, bias, allreduce_dgrad, tp_group): + def forward(ctx, input, weight, bias, allreduce_dgrad, tp_group, output_dtype): """Forward with frozen weight.""" ctx.save_for_backward(weight) ctx.allreduce_dgrad = allreduce_dgrad ctx.tp_group = tp_group - output = torch.matmul(input, weight.t()) - if bias is not None: - output = output + bias - return output + ctx.input_dtype = input.dtype + return _linear_forward(input, weight, bias, output_dtype) @staticmethod @custom_bwd def backward(ctx, grad_output): """Backward with frozen weight.""" (weight,) = ctx.saved_tensors + grad_output = grad_output.to(ctx.input_dtype) if grad_output.dim() > 2: # Work around PyTorch matmul not folding some size-1 leading dims to mm. # Remove this once https://github.com/pytorch/pytorch/issues/186148 is fixed. @@ -439,7 +438,34 @@ def backward(ctx, grad_output): # All-reduce. Note: here async and sync are effectively the same. torch.distributed.all_reduce(grad_input, group=ctx.tp_group) - return grad_input, None, None, None, None + return grad_input, None, None, None, None, None + + +def _linear_forward( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + output_dtype: Optional[torch.dtype], +) -> torch.Tensor: + """Run a linear GEMM with an optional output dtype distinct from its input dtype.""" + if output_dtype is None or output_dtype == input.dtype: + output = torch.matmul(input, weight.t()) + if bias is not None: + output = output + bias + return output + + # Deferred to avoid a circular import: transformer_engine imports tensor-parallel layers. + from megatron.core.extensions.transformer_engine import te_general_gemm + + if te_general_gemm is None: + raise RuntimeError( + "A mixed-precision linear output requires Transformer Engine general_gemm." + ) + + input_shape = input.shape + input_2d = input.reshape(-1, input_shape[-1]) + output = te_general_gemm(weight, input_2d, out_dtype=output_dtype, layout="TN", bias=bias)[0] + return output.reshape(*input_shape[:-1], weight.size(0)) def linear_with_frozen_weight( @@ -453,6 +479,7 @@ def linear_with_frozen_weight( grad_output_buffer: Optional[List[torch.Tensor]] = None, wgrad_deferral_limit: None = None, gtp_remat_size: int = 1, + output_dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: """Linear layer execution with weight.requires_grad == False. @@ -493,6 +520,9 @@ def linear_with_frozen_weight( gtp_remat_size (int): GTP shard count. When > 1 the weight is GTP-sharded and must be all-gathered to its full shape before the matmul, mirroring the trainable path. Defaults to 1 (no-op) for the common non-GTP / non-sharded case. + + output_dtype (torch.dtype optional): Optional GEMM output dtype. A dtype different from + the input dtype requires Transformer Engine ``general_gemm``. """ assert grad_output_buffer is None, ( @@ -516,7 +546,7 @@ def linear_with_frozen_weight( if gtp_remat_size > 1: weight = weight.all_gather_and_prefetch(fwd=True) - args = [input, weight, bias, allreduce_dgrad, tp_group] + args = [input, weight, bias, allreduce_dgrad, tp_group, output_dtype] return LinearWithFrozenWeight.apply(*args) @@ -538,6 +568,7 @@ def forward( wgrad_deferral_limit, tp_group, gtp_remat_size, + output_dtype, ): """Forward.""" if gradient_accumulation_fusion and hasattr(weight, "main_grad"): @@ -561,6 +592,7 @@ def forward( ctx.grad_output_buffer = grad_output_buffer ctx.tp_group = tp_group ctx.gtp_remat_size = gtp_remat_size + ctx.input_dtype = input.dtype if sequence_parallel: dim_size = list(input.size()) @@ -572,10 +604,7 @@ def forward( else: total_input = input - output = torch.matmul(total_input, weight.t()) - if bias is not None: - output = output + bias - return output + return _linear_forward(total_input, weight, bias, output_dtype) @staticmethod @custom_bwd @@ -584,6 +613,9 @@ def backward(ctx, grad_output): input, weight = ctx.saved_tensors main_grad = ctx.main_grad use_bias = ctx.use_bias + # TE owns only the forward GEMM here; Megatron retains the backward contract. + # Cast dY to the input dtype to match the legacy FP32-logit-cast backward path. + grad_output = grad_output.to(ctx.input_dtype) # GTP: re-gather weight for dgrad if ctx.gtp_remat_size > 1: @@ -743,12 +775,13 @@ def backward(ctx, grad_output): None, None, None, + None, ) if ctx.allreduce_dgrad: handle.wait() - return grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None + return grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None, None def linear_with_grad_accumulation_and_async_allreduce( @@ -762,6 +795,7 @@ def linear_with_grad_accumulation_and_async_allreduce( wgrad_deferral_limit: Optional[int] = 0, tp_group: Optional[torch.distributed.ProcessGroup] = None, gtp_remat_size: int = 1, + output_dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: """Linear layer execution with asynchronous communication and gradient accumulation fusion in backprop. @@ -824,6 +858,9 @@ def linear_with_grad_accumulation_and_async_allreduce( wgrad_deferral_limit (int optional): Limit on the number of micro-batches for which embedding weight gradient GEMM should be deferred. Disable by setting this to 0. Defaults to 0. + + output_dtype (torch.dtype optional): Optional GEMM output dtype. A dtype different from + the input dtype requires Transformer Engine ``general_gemm``. """ tp_group = get_tensor_model_parallel_group_if_none(tp_group) @@ -839,6 +876,7 @@ def linear_with_grad_accumulation_and_async_allreduce( wgrad_deferral_limit, tp_group, gtp_remat_size, + output_dtype, ] if not linear_with_grad_accumulation_and_async_allreduce.warned: @@ -912,6 +950,12 @@ class ColumnParallelLinear(torch.nn.Module): If True, reduction of output gradients across tensor-parallel ranks will be disabled. Defaults to False. This feature is used by Lora Adapter in Nemo to delay and fuse reduction along with other gradients for performance optimization. + output_dtype: + Optional dtype for the GEMM output. When it differs from the input dtype, + Transformer Engine ``general_gemm`` is used. + pg_collection: + Optional process group collection. Used to resolve the generalized tensor + parallel remat group; falls back to the global parallel state when omitted. """ def __init__( @@ -934,6 +978,7 @@ def __init__( disable_grad_reduce: bool = False, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + output_dtype: Optional[torch.dtype] = None, pg_collection: Optional[ProcessGroupCollection] = None, ): super(ColumnParallelLinear, self).__init__() @@ -951,6 +996,7 @@ def __init__( self.config = config self.disable_grad_reduce = disable_grad_reduce self.tp_group = tp_group + self.output_dtype = output_dtype self.tp_group = get_tensor_model_parallel_group_if_none( self.tp_group, is_expert=self.is_expert @@ -1180,6 +1226,7 @@ def forward( ), tp_group=self.tp_group, gtp_remat_size=self.gtp_remat_size, + output_dtype=self.output_dtype, ) gather_output = self.gather_output diff --git a/megatron/core/transformer/attention_layer_config.py b/megatron/core/transformer/attention_layer_config.py new file mode 100644 index 00000000000..09b46aaa3f2 --- /dev/null +++ b/megatron/core/transformer/attention_layer_config.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.transformer.transformer_config import TransformerConfig + + +class AttentionLayerConfig(TransformerConfig): + """Configuration for an attention layer in a hybrid stack. + + Due to backwards-compatibility, this config's arguments are defined in TransformerConfig. + """ diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_layer_config.py b/megatron/core/transformer/experimental_attention_variant/dsa_layer_config.py new file mode 100644 index 00000000000..eeab6e9ba3b --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/dsa_layer_config.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.transformer.transformer_config import MLATransformerConfig + + +class DSALayerConfig(MLATransformerConfig): + """Configuration for a DeepSeek Sparse Attention layer in a hybrid stack. + + Due to backwards-compatibility, this config's arguments are defined in MLATransformerConfig. + """ diff --git a/megatron/core/transformer/mla_layer_config.py b/megatron/core/transformer/mla_layer_config.py new file mode 100644 index 00000000000..6a688e01e57 --- /dev/null +++ b/megatron/core/transformer/mla_layer_config.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.transformer.transformer_config import MLATransformerConfig + + +class MLALayerConfig(MLATransformerConfig): + """Configuration for a Multi-Latent Attention layer in a hybrid stack. + + Due to backwards-compatibility, this config's arguments are defined in MLATransformerConfig. + """ diff --git a/megatron/core/transformer/moe/moe_layer_config.py b/megatron/core/transformer/moe/moe_layer_config.py new file mode 100644 index 00000000000..6d3aa8911c4 --- /dev/null +++ b/megatron/core/transformer/moe/moe_layer_config.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.transformer.transformer_config import TransformerConfig + + +class MoELayerConfig(TransformerConfig): + """Configuration for a Mixture-of-Experts layer in a hybrid stack. + + Due to backwards-compatibility, this config's arguments are defined in TransformerConfig. + """ diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 5c55f2abe6c..0f7be974418 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1560,8 +1560,6 @@ def transition_cudagraph_scope(self, mode): and "moe" in self.config.recompute_modules and self.config.cuda_graph_impl == "local" ) - if not hasattr(self, '_router_dtoh_event'): - self._router_dtoh_event = torch.cuda.Event() if not hasattr(self, 'cudagraph_manager_router'): self.cudagraph_manager_router = CudaGraphManager( self.config, self, function_name="_forward_mlp_router" @@ -1630,6 +1628,16 @@ def _get_token_dispatcher_attrs(self): return tuple(attr_names), token_dispatcher_attr_outputs + def _synchronize_router_host_outputs(self, attr_outputs): + """Wait for partial-router graph outputs only when they reside on the host.""" + if not any(attr.device.type == "cpu" for attr in attr_outputs): + return + + if not hasattr(self, '_router_dtoh_event'): + self._router_dtoh_event = torch.cuda.Event() + self._router_dtoh_event.record() + self._router_dtoh_event.synchronize() + def _forward_mlp_router(self, hidden_states, padding_mask=None): """ Executes the router phase of the MoE block. @@ -1735,12 +1743,9 @@ def _forward_mlp_partial_cudagraphs( *token_dispatcher_attr_outputs, ) = router_outputs - # After the router graph replays, the captured .copy_() operations that update - # the returned dispatcher tensors via `_maybe_dtoh_and_synchronize` are queued on - # the current stream but may not have completed. Record an event after the router - # graph and wait on it, so we block only until the router's D2H copies complete. - self._router_dtoh_event.record() - self._router_dtoh_event.synchronize() + # CUDA outputs remain ordered by the graph-completion event. Only host outputs need + # a CPU-blocking wait before the eager dispatcher can consume them. + self._synchronize_router_host_outputs(token_dispatcher_attr_outputs) expert_output, mlp_bias = self._forward_mlp_expert_compute( hidden_states, probs, token_dispatcher_attr_outputs diff --git a/megatron/core/transformer/utils.py b/megatron/core/transformer/utils.py index aee4e961b9e..9983f2f6dc0 100644 --- a/megatron/core/transformer/utils.py +++ b/megatron/core/transformer/utils.py @@ -1,6 +1,7 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. """Utilities for transformer layers.""" + import gc import logging from operator import itemgetter @@ -218,7 +219,7 @@ def make_sharded_object_for_checkpoint( def _get_extra_state_offsets( - sharded_offsets: Iterable[Tuple[int, int, int]] + sharded_offsets: Iterable[Tuple[int, int, int]], ) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: """Turns ShardedTensor offsets into offsets suitable for ShardedObject.""" if sharded_offsets: @@ -300,6 +301,41 @@ def sharded_state_dict_default( _sequence_parallel_attr_cache = None +def set_model_config_attribute(model: Any, attribute: str, value: Any) -> None: + """Set a config attribute on a model and all distinct child-module configs. + + Some models give individual layers separate config objects. Runtime model-wide + toggles must update those configs just as they did when every layer shared the + model's root config. + + Args: + model: Model whose configs should be updated. + attribute: Config attribute to set. + value: Value to assign. The same value object is assigned to every config. + """ + root_config = model.config + setattr(root_config, attribute, value) + updated_config_ids = {id(root_config)} + + module_root = model + visited_wrapper_ids = set() + while not isinstance(module_root, torch.nn.Module) or not hasattr(module_root, "_modules"): + visited_wrapper_ids.add(id(module_root)) + module_root = getattr(module_root, "module", None) + if module_root is None or id(module_root) in visited_wrapper_ids: + return + + for module in module_root.modules(): + config = getattr(module, "config", None) + if ( + config is not None + and id(config) not in updated_config_ids + and hasattr(config, attribute) + ): + setattr(config, attribute, value) + updated_config_ids.add(id(config)) + + def _init_sequence_parallel_cache(model, exclude_modules): """ Initialize the cache of modules with sequence parallel attributes. diff --git a/megatron/elastification/pretrain_hybrid_flex.py b/megatron/elastification/pretrain_hybrid_flex.py index 13eeca1f7c8..adddab019f7 100644 --- a/megatron/elastification/pretrain_hybrid_flex.py +++ b/megatron/elastification/pretrain_hybrid_flex.py @@ -134,6 +134,7 @@ def model_provider(pre_process=True, post_process=True, vp_stage: Optional[int] hybrid_layer_pattern=args.hybrid_layer_pattern, post_process=post_process, fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + logit_dtype=getattr(args, 'logit_dtype', None), parallel_output=True, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, diff --git a/megatron/post_training/model_builder.py b/megatron/post_training/model_builder.py index cd497907406..39c2be07e02 100644 --- a/megatron/post_training/model_builder.py +++ b/megatron/post_training/model_builder.py @@ -383,6 +383,7 @@ def modelopt_gpt_hybrid_builder( "pre_process": pre_process, "post_process": post_process, "fp16_lm_cross_entropy": args.fp16_lm_cross_entropy, + "logit_dtype": getattr(args, "logit_dtype", None), "parallel_output": True, "share_embeddings_and_output_weights": not args.untie_embeddings_and_output_weights, "position_embedding_type": args.position_embedding_type, @@ -428,6 +429,7 @@ def modelopt_gpt_hybrid_builder( "pre_process": pre_process, "post_process": post_process, "fp16_lm_cross_entropy": args.fp16_lm_cross_entropy, + "logit_dtype": getattr(args, "logit_dtype", None), "parallel_output": True, "share_embeddings_and_output_weights": not args.untie_embeddings_and_output_weights, "position_embedding_type": args.position_embedding_type, diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 0b834def67b..e43e2f72ced 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -48,7 +48,11 @@ is_batch_invariant_mode_enabled, ) from megatron.core.transformer.enums import CudaGraphModule -from megatron.core.transformer.utils import toggle_cuda_graphs, transition_moe_cudagraphs +from megatron.core.transformer.utils import ( + set_model_config_attribute, + toggle_cuda_graphs, + transition_moe_cudagraphs, +) from megatron.core.utils import ( get_asyncio_loop, get_attr_wrapped_model, @@ -805,7 +809,7 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa # This is a hack to fix megatron's behaviour when flash-decode affects the training code flow. flash_decode = model.config.flash_decode - model.config.flash_decode = False + set_model_config_attribute(model, "flash_decode", False) fp32_output = not (args.fp16 or args.bf16) with torch.no_grad() if no_grad else nullcontext(): logits_or_hidden_states = model( @@ -816,7 +820,7 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa runtime_gather_output=True, fp32_output=fp32_output, ) - model.config.flash_decode = flash_decode + set_model_config_attribute(model, "flash_decode", flash_decode) pg_collection = get_attr_wrapped_model(model, "pg_collection") pp_group = pg_collection.pp @@ -1704,16 +1708,19 @@ def prepare_data_for_update( data_loader = DataLoader(dataset, batch_size=1) logprobs_batch_size = 1 else: - # Always compute standard masks for the original data (we'll need them later) + # Compute the loss mask and position ids for the original data (we'll need them later). + # No dense attention mask: the forward pass masks via PackedSeqParams (see + # get_logprobs), even when sequence packing is disabled. with nvtx_range("rl/get-ltor-masks", time=True): _, original_loss_mask, original_position_ids = get_ltor_masks_and_position_ids( trajs, tokenizer.eod, tokenizer.pad, args.reset_position_ids, - args.reset_attention_mask, + reset_attention_mask=False, eod_mask_loss=False, pad_mask_loss=True, + create_attention_mask=False, ) original_loss_mask[~generation_masks] = 0.0 compute_trajs = trajs @@ -2200,9 +2207,11 @@ def megatron_rl_inference_mode( # Use local CUDA graphs during rollout inference. An empty module list preserves # full-layer capture when the configured inference scope is layer. - model[0].config.cuda_graph_modules = [] - model[0].config.cuda_graph_impl = "local" - model[0].config.inference_cuda_graph_scope = args.inference_cuda_graph_scope + set_model_config_attribute(model[0], "cuda_graph_modules", []) + set_model_config_attribute(model[0], "cuda_graph_impl", "local") + set_model_config_attribute( + model[0], "inference_cuda_graph_scope", args.inference_cuda_graph_scope + ) # If we get a lower precision wrapper, we go one object deeper. lang_module = model[0].module.module if hasattr(model[0].module, "module") else model[0].module @@ -2260,17 +2269,25 @@ def megatron_rl_inference_mode( # Restore cudagraph scope for training. # MoE partial capture requires specific scopes that aren't user-facing. - model[0].config.cuda_graph_impl = args.cuda_graph_impl - model[0].config.inference_cuda_graph_scope = args.inference_cuda_graph_scope + set_model_config_attribute(model[0], "cuda_graph_impl", args.cuda_graph_impl) + set_model_config_attribute( + model[0], "inference_cuda_graph_scope", args.inference_cuda_graph_scope + ) if args.num_experts is not None: - model[0].config.cuda_graph_modules = [ - CudaGraphModule.mamba, - CudaGraphModule.attn, - CudaGraphModule.moe_router, - CudaGraphModule.moe_preprocess, - ] + set_model_config_attribute( + model[0], + "cuda_graph_modules", + [ + CudaGraphModule.mamba, + CudaGraphModule.attn, + CudaGraphModule.moe_router, + CudaGraphModule.moe_preprocess, + ], + ) else: - model[0].config.cuda_graph_modules = copy.copy(args.cuda_graph_modules) + set_model_config_attribute( + model[0], "cuda_graph_modules", copy.copy(args.cuda_graph_modules) + ) # Switch MoE layers to partial CUDA graph capture for training if args.rl_training_cuda_graphs and args.num_experts is not None: diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 70f26c64d56..124083ead83 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -457,6 +457,7 @@ def gpt_config_from_args( kwargs["transformer_layer_spec"] = import_module(args.spec) kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy + kwargs["logit_dtype"] = getattr(args, "logit_dtype", None) kwargs["position_embedding_type"] = args.position_embedding_type kwargs["rotary_percent"] = args.rotary_percent kwargs["rotary_base"] = args.rotary_base @@ -507,6 +508,7 @@ def hybrid_config_from_args( kwargs["hybrid_stack_spec"] = import_module(args.spec) kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy + kwargs["logit_dtype"] = getattr(args, "logit_dtype", None) kwargs["hybrid_layer_pattern"] = args.hybrid_layer_pattern kwargs["position_embedding_type"] = args.position_embedding_type kwargs["rotary_percent"] = args.rotary_percent diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e54f8a22252..3497834ef93 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1052,6 +1052,7 @@ def validate_args(args, defaults={}): args.mamba_inference_conv_states_dtype = map_dtype(args.mamba_inference_conv_states_dtype) args.mamba_inference_ssm_states_dtype = map_dtype(args.mamba_inference_ssm_states_dtype) args.mamba_training_ssm_states_dtype = map_dtype(args.mamba_training_ssm_states_dtype) + args.logit_dtype = map_dtype(getattr(args, 'logit_dtype', None)) args.megatron_fsdp_main_params_dtype = map_dtype(args.megatron_fsdp_main_params_dtype) args.megatron_fsdp_main_grads_dtype = map_dtype(args.megatron_fsdp_main_grads_dtype) @@ -2935,6 +2936,11 @@ def _add_mixed_precision_args(parser): group.add_argument('--fp16-lm-cross-entropy', action='store_true', help='Move the cross entropy unreduced loss calculation' 'for lm head to fp16.') + group.add_argument('--output-logit-dtype', type=str, choices=['bf16', 'fp32'], default=None, + dest='logit_dtype', + help='Output dtype for the language-model output-layer GEMM. When the ' + 'requested dtype differs from the input dtype, Transformer Engine ' + 'general_gemm is used. By default, logits use the output-layer input dtype.') group.add_argument('--reuse-grad-buf-for-mxfp8-param-ag', action='store_true', help='If True, reuse the grad buffer for MXFP8 parameter all-gather.') group.add_argument('--mamba-training-ssm-states-dtype', type=str, diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index fa859664a46..a81a6c2dfe1 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1748,7 +1748,12 @@ def _load_global_dist_base_checkpoint( ) load_strategy = FullyParallelLoadStrategyWrapper( - load_strategy, process_group, exchange_algo=args.ckpt_fully_parallel_load_exchange_algo + load_strategy, + process_group, + exchange_algo=args.ckpt_fully_parallel_load_exchange_algo, + per_rank_object_load=getattr( + args, 'ckpt_fully_parallel_load_per_rank_objects', False + ), ) if checkpointing_context is not None: checkpointing_context['load_strategy'] = load_strategy diff --git a/megatron/training/config/training_config.py b/megatron/training/config/training_config.py index 1ba743d5962..a379039ff20 100644 --- a/megatron/training/config/training_config.py +++ b/megatron/training/config/training_config.py @@ -548,6 +548,14 @@ class CheckpointConfig: "gather_object": Gather the checkpoint from all ranks in a single operation. """ + ckpt_fully_parallel_load_per_rank_objects: bool = False + """Load ShardedObjects per-rank during fully parallel load of distributed checkpoints. + When True, every rank reads all of its own ShardedObjects (RNG states, + TE `_extra_state`, ...) directly from storage, which removes the WORLD-wide + `all_gather_object` that otherwise exchanges them. Objects are + content-addressable by `unique_key`, so the loaded values are identical. + When False (default), the legacy gather-based object exchange is used.""" + ckpt_fully_parallel_save_process_group: Literal["dp", "ep_dp"] = "dp" """Process group for fully parallel save of distributed checkpoints. "dp"(default): Data parallel process group. diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 87d6aa65b03..c3eb2cccb2e 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -580,3 +580,13 @@ def setup_logging() -> None: if is_rank0(): logger.info(f'Setting logging level to {logging_level}') logging.getLogger().setLevel(logging_level) + + if not is_rank0(): + for noisy_logger_name in [ + 'GroupedGemmQuantSm100', + 'GroupedGemmDsreluSm100', + 'GroupedGemmSreluSm100', + 'GroupedGemmWgradSm100', + 'absl', + ]: + logging.getLogger(noisy_logger_name).setLevel(logging.ERROR) diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py index 63448b196c0..46dcc9b28f4 100644 --- a/megatron/training/models/gpt.py +++ b/megatron/training/models/gpt.py @@ -169,6 +169,7 @@ class GPTModelConfig(ModelConfig): ### GPT Model initialization ### seq_length: int = 1024 fp16_lm_cross_entropy: bool = False + logit_dtype: torch.dtype | None = None parallel_output: bool = True share_embeddings_and_output_weights: bool = False position_embedding_type: Literal["learned_absolute", "rope", "mrope", "yarn", "none"] = "learned_absolute" @@ -316,6 +317,7 @@ def build_model( vocab_size=padded_vocab_size, max_sequence_length=self._model_config.seq_length, fp16_lm_cross_entropy=self._model_config.fp16_lm_cross_entropy, + logit_dtype=self._model_config.logit_dtype, parallel_output=self._model_config.parallel_output, share_embeddings_and_output_weights=self._model_config.share_embeddings_and_output_weights, position_embedding_type=self._model_config.position_embedding_type, diff --git a/megatron/training/models/hybrid.py b/megatron/training/models/hybrid.py index 99f98920eff..c7d1c91b118 100644 --- a/megatron/training/models/hybrid.py +++ b/megatron/training/models/hybrid.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import Any, Callable, ClassVar, Literal, override +import torch + from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.enums import ModelType from megatron.core.models.hybrid.hybrid_layer_specs import ( @@ -50,6 +52,7 @@ class HybridModelConfig(ModelConfig): builder: ClassVar[str] = "megatron.training.models.hybrid.HybridModelBuilder" transformer: TransformerConfig fp16_lm_cross_entropy: bool = False + logit_dtype: torch.dtype | None = None parallel_output: bool = True share_embeddings_and_output_weights: bool = False hybrid_attention_ratio: float = 0.0 @@ -179,6 +182,7 @@ def build_model( max_sequence_length=self._model_config.seq_length, hybrid_layer_pattern=self._model_config.hybrid_layer_pattern, fp16_lm_cross_entropy=self._model_config.fp16_lm_cross_entropy, + logit_dtype=self._model_config.logit_dtype, parallel_output=self._model_config.parallel_output, share_embeddings_and_output_weights=self._model_config.share_embeddings_and_output_weights, position_embedding_type=self._model_config.position_embedding_type, diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py index 30617ef9b4c..1b20eb5a985 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -377,20 +377,44 @@ def get_ltor_masks_and_position_ids(data, reset_position_ids, reset_attention_mask, eod_mask_loss, - pad_mask_loss): - """Build masks and position id for left to right model.""" + pad_mask_loss, + create_attention_mask=True): + """Build masks and position id for left to right model. + + Args: + data: Token ids, shape [micro_batch_size, seq_length]. + eod_token: End-of-document token id. + pad_token: Padding token id. + reset_position_ids: Restart position ids from 0 after each EOD token. + reset_attention_mask: Additionally mask attention across document boundaries, + turning the shared causal mask into a per-sample block-causal mask. + Requires create_attention_mask, since it modifies the materialized mask. + eod_mask_loss: Zero the loss mask at EOD tokens. + pad_mask_loss: Zero the loss mask at pad tokens. + create_attention_mask: Materialize the dense causal attention mask. + Can be disabled if the attention kernel generates the mask by itself + (e.g. from PackedSeqParams), in which case attention_mask is returned as None. + + Returns: + Tuple of (attention_mask or None, loss_mask, position_ids). + """ + assert create_attention_mask or not reset_attention_mask, \ + "reset_attention_mask requires the attention mask to be created." # Extract batch size and sequence length. micro_batch_size, seq_length = data.size() # Attention mask (lower triangular). - if reset_attention_mask: - att_mask_batch = micro_batch_size + if create_attention_mask: + if reset_attention_mask: + att_mask_batch = micro_batch_size + else: + att_mask_batch = 1 + attention_mask = torch.tril( + torch.ones((att_mask_batch, seq_length, seq_length), device=data.device) + ).view(att_mask_batch, 1, seq_length, seq_length) else: - att_mask_batch = 1 - attention_mask = torch.tril( - torch.ones((att_mask_batch, seq_length, seq_length), device=data.device) - ).view(att_mask_batch, 1, seq_length, seq_length) + attention_mask = None # Loss mask. loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device) @@ -429,7 +453,8 @@ def get_ltor_masks_and_position_ids(data, prev_index = i + 1 # Convert attention mask to binary: - attention_mask = attention_mask < 0.5 + if create_attention_mask: + attention_mask = attention_mask < 0.5 return attention_mask, loss_mask, position_ids diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 47c1935eb90..8b979d16d4a 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -17,6 +17,19 @@ if rank != 0: warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) + warnings.filterwarnings("ignore", category=DeprecationWarning) + + # Some libraries (e.g., CUTLASS DSL) use warnings.catch_warnings() with + # simplefilter("always"), which overrides the filters above. Override + # showwarning as a fallback to suppress warnings that slip through. + _original_showwarning = warnings.showwarning + + def _rank0_only_showwarning(message, category, filename, lineno, file=None, line=None): + if issubclass(category, (UserWarning, FutureWarning, DeprecationWarning)): + return + _original_showwarning(message, category, filename, lineno, file, line) + + warnings.showwarning = _rank0_only_showwarning from functools import lru_cache, partial from typing import Any, List, Optional, Tuple diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 39bc7f30b57..cf665b36887 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -16,6 +16,19 @@ if rank != 0: warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) + warnings.filterwarnings("ignore", category=DeprecationWarning) + + # Some libraries (e.g., CUTLASS DSL) use warnings.catch_warnings() with + # simplefilter("always"), which overrides the filters above. Override + # showwarning as a fallback to suppress warnings that slip through. + _original_showwarning = warnings.showwarning + + def _rank0_only_showwarning(message, category, filename, lineno, file=None, line=None): + if issubclass(category, (UserWarning, FutureWarning, DeprecationWarning)): + return + _original_showwarning(message, category, filename, lineno, file, line) + + warnings.showwarning = _rank0_only_showwarning from functools import lru_cache, partial from typing import Any, List, Optional, Tuple diff --git a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_no_load_optim_tp1_pp1_cp1_ep4_dgx_gb200_1N4G/model_config.yaml b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_no_load_optim_tp1_pp1_cp1_ep4_dgx_gb200_1N4G/model_config.yaml index 10d94d27f03..f913c1a62ca 100644 --- a/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_no_load_optim_tp1_pp1_cp1_ep4_dgx_gb200_1N4G/model_config.yaml +++ b/tests/functional_tests/test_cases/nemotron/nemotron3_5_lightning_no_load_optim_tp1_pp1_cp1_ep4_dgx_gb200_1N4G/model_config.yaml @@ -141,6 +141,17 @@ MODEL_ARGS: --ckpt-fully-parallel-load: true --dist-ckpt-strictness: log_all + # Checkpointing optimizations exercised by this load. All three are meant to be + # numerically neutral, so the golden values must stay identical with them on. + # Every rank reads its own ShardedObjects instead of exchanging them with a + # WORLD all_gather_object. + --ckpt-fully-parallel-load-per-rank-objects: true + # Skips the determine_global_metadata all_gather_object that validates sharding. + --no-ckpt-load-validate-sharding-integrity: true + # Redundant TE _extra_state is neither written nor requested on load; only + # delayed-scaling FP8 extra state is ever needed, and this model does not use it. + --ckpt-drop-redundant-extra-state: true + # Validation and functional metrics. eval-interval < train-iters so the # evaluation loop actually runs during the test. --eval-interval: 10 diff --git a/tests/test_utils/python_scripts/resolve_oncall_assignee.py b/tests/test_utils/python_scripts/resolve_oncall_assignee.py new file mode 100644 index 00000000000..8cebeb09b7a --- /dev/null +++ b/tests/test_utils/python_scripts/resolve_oncall_assignee.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Resolve the scheduled Megatron-LM on-call user to an NVIDIA email address.""" + +import argparse +import json +import sys +from contextlib import redirect_stdout +from io import StringIO +from pathlib import Path + +_GITHUB_SCRIPTS_DIR = Path(__file__).resolve().parents[3] / ".github" / "scripts" +sys.path.insert(0, str(_GITHUB_SCRIPTS_DIR)) + +from github_slack_utils import get_user_email # noqa: E402 + + +class AssigneeResolutionError(RuntimeError): + """Raised when the current on-call Linear assignee cannot be resolved safely.""" + + +def load_schedule(path: Path) -> list[dict[str, object]]: + """Load and validate the top-level on-call schedule structure.""" + + try: + schedule = json.loads(path.read_text()) + except OSError as exc: + raise AssigneeResolutionError(f"could not read on-call schedule {path}: {exc}") from exc + except json.JSONDecodeError as exc: + raise AssigneeResolutionError(f"on-call schedule {path} is not valid JSON") from exc + + if not isinstance(schedule, list) or not schedule: + raise AssigneeResolutionError(f"on-call schedule {path} must be a non-empty list") + return schedule + + +def current_oncall(schedule: list[dict[str, object]]) -> str: + """Return the current on-call user, matching the existing rotation manager.""" + + if not schedule: + raise AssigneeResolutionError("on-call schedule must be a non-empty list") + entry = schedule[0] + if not isinstance(entry, dict): + raise AssigneeResolutionError("the current on-call schedule entry must be an object") + + username = entry.get("user") + if not isinstance(username, str) or not username.strip(): + raise AssigneeResolutionError("the current on-call schedule entry must have a user") + return username.strip() + + +def resolve_nvidia_email(github_login: str) -> str: + """Resolve a GitHub login with the repository's shared email lookup helper.""" + + # The shared helper reports lookup details on stdout. Suppress those messages + # so command substitution receives only the email printed by ``main`` and CI + # logs do not repeat email addresses from fallback diagnostics. + try: + with redirect_stdout(StringIO()): + email = get_user_email(github_login) + except SystemExit as exc: + raise AssigneeResolutionError( + f"could not look up an email for GitHub user {github_login!r}" + ) from exc + + if not isinstance(email, str): + raise AssigneeResolutionError( + f"GitHub user {github_login!r} did not resolve to an NVIDIA email" + ) + email = email.strip() + local_part, separator, domain = email.rpartition("@") + if ( + separator != "@" + or not local_part + or "@" in local_part + or domain.casefold() != "nvidia.com" + or any(character.isspace() for character in email) + ): + raise AssigneeResolutionError( + f"GitHub user {github_login!r} did not resolve to a valid NVIDIA email" + ) + return email + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse command-line arguments.""" + + parser = argparse.ArgumentParser( + description="Resolve the scheduled Megatron-LM on-call user's NVIDIA email" + ) + parser.add_argument( + "--schedule-file", + type=Path, + default=Path(".github/oncall_schedule.json"), + help="path to the dated GitHub on-call schedule", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + """Print only the current on-call user's NVIDIA email to stdout.""" + + args = parse_args(argv) + try: + username = current_oncall(load_schedule(args.schedule_file)) + email = resolve_nvidia_email(username) + except AssigneeResolutionError as exc: + print(f"Error: {exc}", file=sys.stderr) + return 1 + + print(f"Resolved scheduled on-call GitHub user {username!r}.", file=sys.stderr) + print(email) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_utils/test_ci_triage.py b/tests/test_utils/test_ci_triage.py index 626013c73ef..3c3b3b7ed08 100644 --- a/tests/test_utils/test_ci_triage.py +++ b/tests/test_utils/test_ci_triage.py @@ -129,13 +129,21 @@ def test_cerno_hard_cutover_contract(): assert triage.count("cerno-linear") == 3 assert triage.count("cerno-notify") == 2 assert triage.count('--config "${CERNO_CONFIG}"') == 3 - assert "ARG CERNO_COMMIT=5a5fb5360e67f8f09d189871bbc0d768c09c43fa" in dockerfile assert '"cerno @ git+${CI_SERVER_URL}/dl/nemo/cerno.git@${CERNO_COMMIT}"' in dockerfile assert "id=CERNO_TOKEN" in dockerfile assert "/run/secrets/CERNO_TOKEN" in dockerfile assert "--secret id=CERNO_TOKEN,env=PAT" in build_script +def test_linear_write_assigns_new_issues_to_scheduled_oncall(): + triage = yaml.safe_load(Path(".gitlab/stages/06.triage.yml").read_text()) + script = "\n".join(triage["triage:linear_write"]["script"]) + + assert "resolve_oncall_assignee.py" in script + assert "--schedule-file .github/oncall_schedule.json" in script + assert '--assignee "${ONCALL_ASSIGNEE}"' in script + + def test_notification_rules_use_expected_pipeline_sources(): unit = yaml.safe_load(Path(".gitlab/stages/02.test.yml").read_text()) functional = yaml.safe_load(Path(".gitlab/stages/04.functional-tests.yml").read_text()) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 413f4ea9a36..8b8cbcfabbc 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -262,16 +262,16 @@ def test_fully_shard_activation_recompute_reshards_parameters(distributed_setup, # Backward completes each module before recomputing the previous one, so # every module-local phase must be cleared after its matching backward. - assert model._phase is FsdpModule.Phase.RESTING - assert model.fc1._phase is FsdpModule.Phase.RESTING - assert model.fc2._phase is FsdpModule.Phase.RESTING + assert model.phase is FsdpModule.Phase.RESTING + assert model.fc1.phase is FsdpModule.Phase.RESTING + assert model.fc2.phase is FsdpModule.Phase.RESTING # A second forward after backward runs in the forward phase again, so # forward-order prefetch resumes and the module phases return to resting. model(x).sum().backward() - assert model._phase is FsdpModule.Phase.RESTING - assert model.fc1._phase is FsdpModule.Phase.RESTING - assert model.fc2._phase is FsdpModule.Phase.RESTING + assert model.phase is FsdpModule.Phase.RESTING + assert model.fc1.phase is FsdpModule.Phase.RESTING + assert model.fc2.phase is FsdpModule.Phase.RESTING @pytest.mark.parametrize("set_to_none", [True, False]) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py index a132ff88139..3892a354540 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py @@ -7,6 +7,7 @@ import pytest import torch +import megatron.core.distributed.fsdp.mcore_fsdp_adapter as mcore_fsdp_adapter from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module import FsdpModule @@ -73,7 +74,6 @@ def test_wraps_fsdp_unit_modules_before_root(self): data_parallel_sharding_strategy="optim_grads_params", megatron_fsdp_main_params_dtype=torch.float32, megatron_fsdp_main_grads_dtype=torch.float32, - fsdp_all_gather_in_start_param_sync=False, ), module=model, fsdp_unit_modules=[TransformerLayer], @@ -100,6 +100,40 @@ def test_wraps_fsdp_unit_modules_before_root(self): assert child_parameter_names assert root_parameter_names == {"1.weight", "1.bias"} + def test_nccl_ub_enables_symmetric_memory(self, monkeypatch): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + ffn_hidden_size=32, + bf16=True, + params_dtype=torch.bfloat16, + ) + model = torch.nn.Linear(config.hidden_size, config.hidden_size).to( + device="cuda", dtype=config.params_dtype + ) + fully_shard_context_calls = [] + original_fully_shard_context = mcore_fsdp_adapter.fully_shard_context + + def record_fully_shard_context(*args, **kwargs): + fully_shard_context_calls.append(kwargs["use_symmetric_memory"]) + return original_fully_shard_context(*args, **kwargs) + + monkeypatch.setattr(mcore_fsdp_adapter, "fully_shard_context", record_fully_shard_context) + FullyShardedDataParallel( + config=config, + ddp_config=DistributedDataParallelConfig( + use_megatron_fsdp=True, + megatron_fsdp_version=2, + data_parallel_sharding_strategy="optim_grads_params", + nccl_ub=True, + ), + module=model, + pg_collection=self.pg_collection, + ) + + assert fully_shard_context_calls == [True] + def test_build_train_and_step(self): config = TransformerConfig( num_layers=2, @@ -128,7 +162,6 @@ def test_build_train_and_step(self): data_parallel_sharding_strategy="optim_grads_params", megatron_fsdp_main_params_dtype=torch.float32, megatron_fsdp_main_grads_dtype=torch.bfloat16, - fsdp_all_gather_in_start_param_sync=False, ), module=model, pg_collection=self.pg_collection, @@ -164,7 +197,7 @@ def test_build_train_and_step(self): torch.randn(8, 2, config.hidden_size, device="cuda", dtype=torch.bfloat16) for _ in range(2) ] - for _ in range(3) + for _ in range(10) ] reference_losses = [] @@ -198,4 +231,4 @@ def test_build_train_and_step(self): reference_losses = torch.stack(reference_losses) assert torch.isfinite(losses).all() assert torch.isfinite(reference_losses).all() - torch.testing.assert_close(losses, reference_losses, rtol=1e-2, atol=0) + torch.testing.assert_close(losses, reference_losses, rtol=1e-3, atol=0) diff --git a/tests/unit_tests/distributed/test_layer_wise_param_layout.py b/tests/unit_tests/distributed/test_layer_wise_param_layout.py index a73bf4789a5..2d68a96d7ad 100644 --- a/tests/unit_tests/distributed/test_layer_wise_param_layout.py +++ b/tests/unit_tests/distributed/test_layer_wise_param_layout.py @@ -106,6 +106,15 @@ def _assert_param_within_shard(layout, param, dp_size): class TestSizeMatchingLayout: + """Packing and bucketing rules, exercised with 1-D params. + + Every param here is 1-D, so ``_ns_compute_cost`` falls back to ``nelement()`` + and Newton-Schulz cost equals numel. That keeps these cases focused on the + packing and bucketing rules, but it also means they cannot tell + compute-balanced placement apart from numel-balanced placement: both order + and assign identically when the two metrics agree. ``TestComputeBalancedLayout`` + covers that distinction with 2-D GTP-sharded params, where the metrics diverge. + """ # -- uniform params: all same size, dp_size divides count -- @@ -318,6 +327,32 @@ def test_bucket_has_no_padding_when_params_pack_evenly(self): total_buffer_numel = layout.bucket_indices[-1][1] assert total_buffer_numel == 8 * numel + def test_mixed_sizes_can_absorb_into_larger_bucket(self): + """Absorbing is not always a win: with mixed sizes it can cost space. + + ``_place`` walks params in backprop order while ``_emit_bucket`` sorts the + chunk first, so for mixed sizes the two reach different shard loads. + ``_absorbs`` compares against its own lower estimate and admits the two + 128-element params, after which the sorted packing stacks both onto one + shard and the bucket grows. Closing at the threshold instead would have + emitted 768 + 512 = 1280 elements; absorbing emits 1024 + 384 = 1408. + + Documented rather than fixed: the target case is equal-sized expert + matrices, where sorting is a no-op and the estimate is exact. + """ + dp_size = 2 + # Backprop order is reversed(params), so this list is written back to front. + backprop_order_numels = [192, 192, 256, 128, 128, 192] + params = [_make_param((n,)) for n in reversed(backprop_order_numels)] + cfg = _make_ddp_config() + + layout = _LWO._compute_per_buffer_param_layout(params, 448, dp_size, cfg) + + for param in params: + _assert_param_within_shard(layout, param, dp_size) + assert len(layout.bucket_indices) == 2 + assert layout.bucket_indices[-1][1] == 1408 + # -- bucket alignment -- def test_bucket_dp_divisible(self): @@ -427,3 +462,102 @@ def test_expert_parallel_separate_buffer(self): cfg = _make_ddp_config() layout = _LWO.compute_full_param_layout([dense, expert], None, dp_size, cfg) assert len(layout.layouts) == 2 + + +# --------------------------------------------------------------------------- +# Tests for compute-balanced LPT (Newton-Schulz cost, not numel) +# --------------------------------------------------------------------------- + + +class TestComputeBalancedLayout: + """Placement keys on Newton-Schulz cost, so GTP-sharded params spread out. + + A GTP-sharded param has the numel of its local shard but Newton-Schulz runs on + the full all-gathered matrix, so its cost is far higher than numel suggests. + """ + + def _ns_cost(self, param): + rows, cols = param.data.shape + rows *= getattr(param, 'gtp_remat_size', 1) + big, small = max(rows, cols), min(rows, cols) + return big * small * small + + def _shard_compute_loads(self, layout, params, dp_size): + loads = [0] * dp_size + for param in params: + loads[_get_shard_for_param(layout, param, dp_size)] += self._ns_cost(param) + return loads + + def test_gtp_params_balanced_by_compute_not_numel(self): + """GTP-sharded params dominate cost while having the smallest numel. + + Sorting by numel puts the three cheap-but-large params first and leaves the + expensive GTP ones to fill in, which piles them onto shards that are already + loaded. Sorting by compute cost spreads them instead. + """ + dp_size = 4 + gtp = [ + _make_param((64, 256), is_gtp_weight_remat=True, gtp_remat_size=64) for _ in range(3) + ] + dense = [_make_param((128, 1024)), _make_param((256, 256)), _make_param((256, 256))] + tail = [_make_param((64, 256))] + params = dense + gtp + tail + cfg = _make_ddp_config() + + layout = _LWO._compute_per_buffer_param_layout(params, None, dp_size, cfg) + + # Each GTP param costs 268M against 16.8M for the largest dense param, so no + # two of the three may share a shard. + gtp_shards = [_get_shard_for_param(layout, param, dp_size) for param in gtp] + assert len(set(gtp_shards)) == len(gtp), f"GTP params clustered onto {gtp_shards}" + + loads = self._shard_compute_loads(layout, params, dp_size) + imbalance = max(loads) / (sum(loads) / dp_size) + # Numel-ordered placement gives 3.77x on this input. + assert imbalance < 1.5, f"compute imbalance {imbalance:.2f}x too high: {loads}" + + def test_gtp_params_land_on_different_shards_in_each_bucket(self): + """Expensive params spread across buckets because compute loads do not reset. + + Four buckets, each holding one GTP-sharded matrix plus three dense params of + identical numel. ``shard_cursors`` resets per bucket, so numel gives every + bucket the same starting state and would send all four GTP matrices to shard + 0. ``shard_compute_loads`` carries over, so each bucket sees the previous + ones' cost and picks a different shard. + """ + dp_size = 4 + buckets = 4 + params = [] + for _ in range(buckets): + params.append(_make_param((64, 256), is_gtp_weight_remat=True, gtp_remat_size=64)) + params.extend(_make_param((128, 128)) for _ in range(3)) + gtp_params = [param for param in params if hasattr(param, 'gtp_remat_size')] + cfg = _make_ddp_config() + + # Each group of four params is 4 * 16384 elements, so this cuts one bucket per group. + layout = _LWO._compute_per_buffer_param_layout(params, 4 * 16384, dp_size, cfg) + + assert len(layout.bucket_indices) == buckets + gtp_shards = [_get_shard_for_param(layout, param, dp_size) for param in gtp_params] + assert len(set(gtp_shards)) == buckets, f"GTP params clustered onto {gtp_shards}" + + def test_param_larger_than_cap_is_still_placed(self): + """A param larger than the per-bucket cap still gets placed. + + The cap is ``total_chunk_numel / dp_size * 1.3``. A param above it disqualifies + every shard, so assignment falls back to the previous least-numel rule instead + of wedging. + """ + dp_size = 2 + oversized = _make_param((1024,)) + small_params = [_make_param((64,)) for _ in range(2)] + cfg = _make_ddp_config() + + # cap = (1024 + 64 + 64) / 2 * 1.3 = 748.8, below the oversized param's numel. + layout = _LWO._compute_per_buffer_param_layout( + small_params + [oversized], None, dp_size, cfg + ) + + for param in small_params + [oversized]: + _assert_param_within_shard(layout, param, dp_size) + assert oversized in layout.param_index_map diff --git a/tests/unit_tests/extension/test_te_lmhead_column_parallel_linear.py b/tests/unit_tests/extension/test_te_lmhead_column_parallel_linear.py index dbde4c67df4..7d5d58777af 100644 --- a/tests/unit_tests/extension/test_te_lmhead_column_parallel_linear.py +++ b/tests/unit_tests/extension/test_te_lmhead_column_parallel_linear.py @@ -92,6 +92,10 @@ def test_rejects_disable_grad_reduce(self): with pytest.raises(ValueError, match="disable_grad_reduce"): TELMHeadColumnParallelLinear(**self._kwargs(disable_grad_reduce=True)) + def test_rejects_output_dtype(self): + with pytest.raises(ValueError, match="output_dtype"): + TELMHeadColumnParallelLinear(**self._kwargs(output_dtype=torch.float32)) + class TestGPTModelOutputLayerSelection: """Verify GPTModel picks the right output-layer class based on config.""" @@ -118,6 +122,20 @@ def test_default_uses_column_parallel_linear(self): assert isinstance(model.output_layer, tensor_parallel.ColumnParallelLinear) assert not isinstance(model.output_layer, TELMHeadColumnParallelLinear) + @pytest.mark.internal + def test_logit_dtype_is_forwarded_to_output_layer(self): + config = TransformerConfig( + num_layers=2, hidden_size=12, num_attention_heads=4, use_cpu_initialization=True + ) + model = GPTModel( + config=config, + transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(), + vocab_size=100, + max_sequence_length=4, + logit_dtype=torch.float32, + ) + assert model.output_layer.output_dtype == torch.float32 + @pytest.mark.internal @pytest.mark.skipif( not _IS_BLACKWELL, reason="MXFP8 output projection requires Blackwell (SM >= 10)" diff --git a/tests/unit_tests/generalized_tensor_parallel/gtp_test_utils.py b/tests/unit_tests/generalized_tensor_parallel/gtp_test_utils.py index 259cc6ed0d5..dab9fa0790c 100644 --- a/tests/unit_tests/generalized_tensor_parallel/gtp_test_utils.py +++ b/tests/unit_tests/generalized_tensor_parallel/gtp_test_utils.py @@ -1,7 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Shared fixtures and helpers for all GTP unit tests. -""" +"""Shared fixtures and helpers for all GTP unit tests.""" import pytest import torch @@ -9,7 +8,10 @@ from transformer_engine.pytorch import is_mxfp8_available, is_nvfp4_available from transformer_engine.pytorch.quantization import FP8GlobalStateManager -from megatron.core.tensor_parallel.generalized_tensor_parallelism import GTPShardedParam +from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( + GTPShardedParam, + reset_gtp_state, +) from tests.unit_tests.test_utilities import Utils # --------------------------------------------------------------------------- @@ -33,9 +35,16 @@ def reset_fp8_state(): @pytest.fixture(autouse=True) def reset_gtp_globals(): - """Reset GTP mutable class-level state between tests.""" + """Reset GTP mutable class-level state between tests. + + Defers to the production reset so this cannot drift as new class-level state is added. + Note it only clears the process-global cursors: the chain links themselves live on the + params (``prev_w`` / ``_recompute_prev`` / ``_ag_ticket_*``, set once in + ``_init_gtp_runtime_attrs``), so a test that reuses modules across cases would inherit + stale links. Every GTP test builds fresh modules for this reason. + """ yield - GTPShardedParam._chain_state = {} + reset_gtp_state() # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py index 297c89578f1..6c594d0be62 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py @@ -502,6 +502,42 @@ def test_graphness_helpers(self): assert gtp_module._chain_is_graphed("GTP_graphed") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="GTPShardedParam requires CUDA") +class TestLatentProjectionChainClassification: + """Classify opted-in latent projections through the public GTP chain API.""" + + FC1 = "decoder.layers.3.mlp.fc1_latent_proj.weight" + FC2 = "decoder.layers.3.mlp.fc2_latent_proj.weight" + + def teardown_method(self, method): + gtp_module.set_cuda_graph_modules(None, cuda_graph_impl="none") + gtp_module.reset_gtp_state() + + def _chains(self, *, cuda_graph_modules=None, cuda_graph_impl="none"): + params = tuple(GTPShardedParam(torch.zeros(1, device="cuda")) for _ in range(2)) + assert all(isinstance(param, GTPShardedParam) for param in params) + + class _Model: + def named_parameters(_self): + return iter(zip((self.FC1, self.FC2), params)) + + gtp_module.classify_gtp_remat_chains( + _Model(), cuda_graph_modules=cuda_graph_modules, cuda_graph_impl=cuda_graph_impl + ) + return tuple(param.chain_id for param in params) + + def test_eager_latent_projections_are_ungraphed(self): + assert self._chains() == (GTPChain.UNGRAPHED.value, GTPChain.UNGRAPHED.value) + + def test_local_router_captures_both_latent_projections(self): + chains = self._chains(cuda_graph_modules={"moe_router"}, cuda_graph_impl="local") + assert chains == (GTPChain.GRAPHED.value, GTPChain.GRAPHED.value) + + def test_unrelated_local_scope_leaves_latent_projections_ungraphed(self): + chains = self._chains(cuda_graph_modules={"mamba", "attn"}, cuda_graph_impl="local") + assert chains == (GTPChain.UNGRAPHED.value, GTPChain.UNGRAPHED.value) + + class TestGroupedDoubleBuffer: """One-block-ahead grouped chains must double-buffer: consecutive MoE layers get distinct gather buffers (else prefetching layer N+1 clobbers layer N's in-use weight). Pure cache-key diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py index addcb9cb657..14641c3a3b4 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py @@ -36,9 +36,10 @@ ) -def _worker_gtp_partial_cg_correctness(rank, world_size, port): - """Compare eager and local attention CUDA graphs with GTP2 x DP2.""" +def _worker_gtp_partial_cg_correctness(rank, world_size, port, partial_cg_modules, opt_in_modules): + """Compare eager and local CUDA graphs with GTP2 x DP2.""" del port + gtp_module._GTP_PARAMS.clear() from megatron.core import parallel_state as ps from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec @@ -56,14 +57,15 @@ def _worker_gtp_partial_cg_correctness(rank, world_size, port): ) from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp from megatron.core.transformer.transformer_config import TransformerConfig - - hidden = 4096 - num_heads = 32 - ffn_hidden = 16384 - # Four layers force parameters with matching scheduling domains/shapes to reuse the two-slot - # wgrad ring across independently replayed graphs. - num_layers = 4 - sequence_length = 32 + from megatron.core.transformer.transformer_layer import MoETransformerLayer + + latent_projection_case = "moe_latent_proj" in opt_in_modules + hidden = 256 if latent_projection_case else 4096 + num_heads = 8 if latent_projection_case else 32 + ffn_hidden = 512 if latent_projection_case else 16384 + # Use multiple layers to exercise repeated local CUDA-graph execution with GTP parameters. + num_layers = 1 if latent_projection_case else 4 + sequence_length = 16 if latent_projection_case else 32 batch_size = 1 learning_rate = 0.01 steps = 10 @@ -73,6 +75,18 @@ def _worker_gtp_partial_cg_correctness(rank, world_size, port): assert world_size == gtp_degree * dp_degree def make_config(*, partial_cg=False): + moe_options = {} + if latent_projection_case: + moe_options = { + "num_moe_experts": 2, + "moe_router_topk": 1, + "moe_router_pre_softmax": True, + "moe_ffn_hidden_size": ffn_hidden, + "moe_grouped_gemm": True, + "moe_token_dispatcher_type": "allgather", + "moe_aux_loss_coeff": 0.0, + "moe_latent_size": 128, + } return TransformerConfig( num_attention_heads=num_heads, num_layers=num_layers, @@ -83,16 +97,40 @@ def make_config(*, partial_cg=False): hidden_dropout=0.0, attention_dropout=0.0, bias_dropout_fusion=False, + gradient_accumulation_fusion=latent_projection_case, tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_weight_remat_size=gtp_degree, + gtp_remat_opt_in_modules=opt_in_modules, cuda_graph_impl="local" if partial_cg else "none", - cuda_graph_modules=["attn"] if partial_cg else [], + cuda_graph_modules=partial_cg_modules if partial_cg else [], cuda_graph_warmup_steps=2, + **moe_options, ) - def make_attention_stack(config, pg_collection): - spec = copy.deepcopy(get_gpt_layer_with_transformer_engine_spec()) + def make_layer_stack(config, pg_collection): + spec = copy.deepcopy( + get_gpt_layer_with_transformer_engine_spec( + num_experts=2 if latent_projection_case else None, + moe_grouped_gemm=latent_projection_case, + ) + ) + if latent_projection_case: + spec.submodules.input_layernorm = IdentityOp + spec.submodules.self_attention = IdentityOp + spec.submodules.self_attn_bda = IdentityFuncOp + return torch.nn.ModuleList( + [ + MoETransformerLayer( + config, + spec.submodules, + layer_number=1, + pg_collection=pg_collection, + name="decoder.layers.0", + ) + ] + ) + spec.submodules.pre_mlp_layernorm = IdentityOp spec.submodules.mlp = IdentityOp spec.submodules.mlp_bda = IdentityFuncOp @@ -105,6 +143,27 @@ def make_attention_stack(config, pg_collection): ] ) + def get_cudagraph_managers(layers): + if latent_projection_case: + return [ + manager + for layer in layers + for manager in (layer.cudagraph_manager_router, layer.cudagraph_manager_postprocess) + ] + return [layer.cudagraph_manager for layer in layers] + + def get_latent_params(layers): + return [ + param + for name, param in layers.named_parameters() + if "fc1_latent_proj.weight" in name or "fc2_latent_proj.weight" in name + ] + + def make_pg_collection(): + if latent_projection_case: + return ProcessGroupCollection.use_mpu_process_groups() + return ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp", "gtp_remat"]) + def run_step(layers, x): with fp8_autocast(enabled=False): for layer in layers: @@ -162,17 +221,19 @@ def apply_sgd_step(layers, gtp_size): tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=gtp_degree ) model_parallel_cuda_manual_seed(42) - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["tp", "cp", "gtp_remat"] - ) + pg_collection = make_pg_collection() eager_config = make_config() - eager = make_attention_stack(eager_config, pg_collection).cuda() + eager = make_layer_stack(eager_config, pg_collection).cuda() eager_gtp_group = ps.get_gtp_weight_remat_group() eager_dp_group = ps.get_data_parallel_group(with_gtp_remat=False) eager_dp_rank = eager_dp_group.rank() assert eager_gtp_group.size() == gtp_degree assert eager_dp_group.size() == dp_degree assert any(isinstance(param, GTPShardedParam) for param in eager.parameters()) + if latent_projection_case: + eager_latent_params = get_latent_params(eager) + assert len(eager_latent_params) == 2 + assert all(isinstance(param, GTPShardedParam) for param in eager_latent_params) initialize_main_grads(eager) saved_local_weights = {name: param.data.clone() for name, param in eager.named_parameters()} @@ -190,20 +251,20 @@ def apply_sgd_step(layers, gtp_size): apply_sgd_step(eager, eager_gtp_group.size()) del eager, loss, x + torch.cuda.synchronize() ps.destroy_model_parallel() gtp_module.reset_gtp_state() + gtp_module._GTP_PARAMS.clear() - # Optimized path: the same GTP2 x DP2 topology with attention-only local CUDA graphs. + # Optimized path: the same GTP2 x DP2 topology with local CUDA graphs. ps.initialize_model_parallel( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=gtp_degree ) initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) model_parallel_cuda_manual_seed(42) - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["tp", "cp", "gtp_remat"] - ) + pg_collection = make_pg_collection() partial_cg_config = make_config(partial_cg=True) - partial_cg = make_attention_stack(partial_cg_config, pg_collection).cuda() + partial_cg = make_layer_stack(partial_cg_config, pg_collection).cuda() classify_gtp_remat_chains( partial_cg, cuda_graph_modules=partial_cg_config.cuda_graph_modules, @@ -220,7 +281,8 @@ def apply_sgd_step(layers, gtp_size): assert dp_rank == eager_dp_rank gtp_params = [param for param in partial_cg.parameters() if isinstance(param, GTPShardedParam)] assert gtp_params, "GTP not active: no GTPShardedParam found" - assert all(param.chain_id == GTPChain.GRAPHED.value for param in gtp_params) + params_in_scope = get_latent_params(partial_cg) if latent_projection_case else gtp_params + assert all(param.chain_id == GTPChain.GRAPHED.value for param in params_in_scope) for name, param in partial_cg.named_parameters(): param.data.copy_(saved_local_weights[name]) # Production captures after DDP maps every parameter into a main-grad buffer and initializes @@ -241,18 +303,21 @@ def apply_sgd_step(layers, gtp_size): wait_for_gtp_grad_reduction_on_current_stream() eager_grad_norm = global_grad_norm(partial_cg, gtp_group) eager_probe_loss_value = eager_probe_loss.item() + del eager_probe_loss, eager_probe_x + reset_grad_state(partial_cg) create_cudagraphs() assert _CudagraphGlobalRecord.cudagraph_created - runners = [layer.cudagraph_manager.cudagraph_runners[0] for layer in partial_cg] - assert all(runner.gtp_remat for runner in runners) - assert any(runner._gtp_wgrad_ring_slots for runner in runners) + managers = get_cudagraph_managers(partial_cg) + assert all(len(manager.cudagraph_runners) == 1 for manager in managers) + runners = [manager.cudagraph_runners[0] for manager in managers] + assert any(runner.gtp_remat for runner in runners) replay_grad_norms = [] replay_losses = [] for _ in range(3): reset_grad_state(partial_cg) - replay_x = eager_probe_x.detach().clone().requires_grad_() + replay_x = make_replica_input(1234, dp_rank).requires_grad_() replay_loss = run_step(partial_cg, replay_x) replay_loss.backward() wait_for_gtp_grad_reduction_on_current_stream() @@ -280,7 +345,7 @@ def apply_sgd_step(layers, gtp_size): flush=True, ) - del eager_probe_loss, eager_probe_x, replay_loss, replay_x + del replay_loss, replay_x for step in range(steps): reset_grad_state(partial_cg) @@ -295,19 +360,21 @@ def apply_sgd_step(layers, gtp_size): del loss, x finally: torch.cuda.synchronize() - for layer in partial_cg: - for runner in layer.cudagraph_manager.cudagraph_runners: + managers = get_cudagraph_managers(partial_cg) + for manager in managers: + for runner in manager.cudagraph_runners: if runner.fwd_graph is not None: runner.fwd_graph.reset() if runner.bwd_graph is not None: runner.bwd_graph.reset() delete_cuda_graphs() - for layer in partial_cg: - layer.cudagraph_manager.cudagraph_runners.clear() + for manager in managers: + manager.cudagraph_runners.clear() gc.collect() ps.destroy_model_parallel() ps.initialize_model_parallel() gtp_module.reset_gtp_state() + gtp_module._GTP_PARAMS.clear() if rank == 0: for step, (eager_loss, partial_cg_loss) in enumerate(zip(eager_losses, partial_cg_losses)): @@ -324,8 +391,17 @@ def apply_sgd_step(layers, gtp_size): class TestGTPPartialCGCorrectness: - def test_gtp_partial_cg_loss_and_grad_norm_match_eager(self): + @pytest.mark.parametrize( + "partial_cg_modules,opt_in_modules", + [ + pytest.param(["attn"], [], id="attention"), + pytest.param(["moe_router"], ["moe_latent_proj"], id="moe-router-latent-projections"), + ], + ) + def test_gtp_partial_cg_loss_and_grad_norm_match_eager( + self, partial_cg_modules, opt_in_modules + ): """Local-CG loss trajectory and global grad norm must match eager execution.""" if torch.cuda.device_count() < 4: pytest.skip("Requires at least 4 CUDA devices") - _run_distributed(_worker_gtp_partial_cg_correctness, 4) + _run_distributed(_worker_gtp_partial_cg_correctness, 4, partial_cg_modules, opt_in_modules) diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_recompute_chain.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_recompute_chain.py new file mode 100644 index 00000000000..66367c9a86a --- /dev/null +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_recompute_chain.py @@ -0,0 +1,295 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Regression tests for the GTP recompute-forward prefetch chain. + +Weights re-gathered during an activation-recompute forward form their own chain and prefetch one +node ahead, so two adjacent nodes sharing a gather buffer is a data race: the prefetch of node +i+1 overwrites the weight node i is still reading. In training that shows up as silently wrong +recomputed activations, exploding grad norm, then NaN. + +The buffer tests are the regression guard. The clobber is a stream race, so a numerical test +only trips when the timing lines up -- removing the parity fails the buffer tests but not the +numerical one at this size. TestGroupedDoubleBuffer asserts cache keys for the same reason. + +Test groups +----------- +TestGTPRecomputeChainBuffers - adjacent recompute nodes never share a gather buffer +TestGTPRecomputeCorrectness - recompute reproduces the non-recompute dgrad and wgrads +TestGroupedGTPRecomputeChainBuffers - same invariant per expert on a grouped chain +""" + +import pytest +import torch +import torch.distributed as dist + +from megatron.core.tensor_parallel.gtp_api import HAVE_GTP + +if not HAVE_GTP: + pytest.skip("GTP requires TransformerEngine >= 2.19", allow_module_level=True) + +import megatron.core.tensor_parallel.generalized_tensor_parallelism as gtp_module +from megatron.core.tensor_parallel.random import CheckpointWithoutOutput +from tests.unit_tests.generalized_tensor_parallel.gtp_test_utils import ( + _make_gtp_linear, + _make_gtp_remat_grouped_linear, + _requires_multi_gpu, + _run_distributed, + _torchrun_dist_init, + reset_fp8_state, + reset_gtp_globals, +) + +# Topology matters. Each block is one recomputed square GEMM followed by two NON-recomputed +# GEMMs of a different shape: +# +# fwd chain: A0 (H,H) -> B0 (O,H) -> C0 (H,O) -> A1 (H,H) -> ... heterogeneous +# recompute chain: A0 -------------------------------> A1 -> ... homogeneous +# +# The differently-shaped B/C are what make this a real test: without them the A weights are +# adjacent in the FORWARD chain too, _ensure_no_shared_buffer_with separates them there, +# and the recompute-chain guard becomes redundant -- so the test would pass even when broken. +HIDDEN = 512 +OTHER = 256 +NUM_LAYERS = 4 +DTYPE = torch.bfloat16 + + +def _recompute_buffer_addr(param): + """Address of the buffer this weight gathers into on the recompute chain, or None. + + Test-only: the production code has no reason to know a buffer's address. + """ + ticket = getattr(param, "_ag_ticket_recompute", None) + if ticket is None: # ticket ids start at 0, so compare against None + return None + slot = gtp_module.get_global_GTP_cache()._slots.get(ticket) + buf = slot.buf if slot is not None else None + if buf is None: + return None + raw = getattr(buf, "_rowwise_data", None) + if raw is None: + raw = getattr(buf, "_data", buf) + return raw.data_ptr() + + +def _build_layers(world_size): + """Return (recomputed square layers, non-recomputed differently-shaped spacer pairs).""" + gtp_remat_group = dist.new_group(list(range(world_size))) + + def linear(in_f, out_f): + return _make_gtp_linear(in_f, out_f, gtp_remat_group, DTYPE, fuse_wgrad_accumulation=True) + + recomputed = [linear(HIDDEN, HIDDEN) for _ in range(NUM_LAYERS)] + spacers = [(linear(HIDDEN, OTHER), linear(OTHER, HIDDEN)) for _ in range(NUM_LAYERS)] + for layer in recomputed + [m for pair in spacers for m in pair]: + # GTP reduce-scatters wgrad into main_grad, on the local shard shape. + layer.weight.main_grad = torch.zeros(layer.weight.shape, dtype=DTYPE, device="cuda") + return recomputed, spacers + + +def _zero_grads(layers): + recomputed, spacers = layers + for layer in recomputed + [m for pair in spacers for m in pair]: + layer.weight.main_grad.zero_() + + +def _forward_backward(layers, x, recompute): + """Run the stack, optionally checkpointing every layer so it is recomputed in backward. + + Mirrors the production pattern: checkpoint the GTP GEMM, let a downstream op consume (and + save) its output, then hook the recompute on that downstream tensor so it fires during this + layer's backward. + """ + recomputed, spacers = layers + h = x + for layer, (spacer_down, spacer_up) in zip(recomputed, spacers): + # te.Linear returns a bare tensor (bias=False), not Megatron's (out, bias) tuple. + if recompute: + checkpoint = CheckpointWithoutOutput() + y = checkpoint.checkpoint(lambda inp, l=layer: l(inp), h) + # gelu saves y for backward, so discarding y is what makes the recompute necessary. + h = torch.nn.functional.gelu(y) + checkpoint.discard_output_and_register_recompute(h) + else: + h = torch.nn.functional.gelu(layer(h)) + # Not checkpointed: keeps the fwd chain heterogeneous around each recomputed weight. + h = torch.nn.functional.gelu(spacer_down(h)) + h = torch.nn.functional.gelu(spacer_up(h)) + loss = h.float().sum() + loss.backward() + # Return the INPUT gradient, not the loss: the loss is produced by the forward pass, so it + # is identical either way and cannot witness anything the recompute got wrong. + return x.grad.detach().clone() + + +def _worker_adjacent_nodes_use_distinct_buffers(rank, world_size, port): + """Every adjacent pair of recompute-chain nodes must gather into different buffers.""" + torch.manual_seed(0) + gtp_module.reset_gtp_state() + layers = _build_layers(world_size) + recomputed = layers[0] + + x = torch.randn(8, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=True) + dist.broadcast(x, src=0) + + # First pass builds the chain (all gathers on demand); the second uses it. + for _ in range(2): + _forward_backward(layers, x, recompute=True) + + chain = [] + node = recomputed[-1].weight + while node is not None and node._recompute_prev is not None: + node = node._recompute_prev + while node is not None: + chain.append(node) + node = node._recompute_next + + assert len(chain) == NUM_LAYERS, f"recompute chain has {len(chain)} nodes, want {NUM_LAYERS}" + + addrs = [_recompute_buffer_addr(w) for w in chain] + assert all(a is not None for a in addrs), f"unallocated recompute buffer: {addrs}" + shared = [i for i in range(len(addrs) - 1) if addrs[i] == addrs[i + 1]] + assert not shared, ( + f"recompute-chain nodes {shared} share a gather buffer with their successor " + f"(addrs={[hex(a) for a in addrs]}); the one-ahead prefetch would clobber the weight " + "still being read" + ) + # One-ahead needs exactly two buffers; more would mean the pool stopped being reused. + assert len(set(addrs)) == 2, f"want 2 alternating buffers, got {len(set(addrs))}: {addrs}" + + +def _worker_recompute_matches_no_recompute(rank, world_size, port): + """Recompute is pure rematerialization: same input grad and same weight grads. + + End-to-end sanity check over the real CheckpointWithoutOutput path. It does not reliably + catch the buffer-sharing bug on its own (see the module docstring); the buffer invariant is + what guards that. + """ + torch.manual_seed(0) + gtp_module.reset_gtp_state() + layers = _build_layers(world_size) + recomputed = layers[0] + + x = torch.randn(8, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=True) + dist.broadcast(x, src=0) + + # Warm up so the chains exist, then measure a steady-state step of each variant. + _forward_backward(layers, x, recompute=False) + _forward_backward(layers, x, recompute=True) + + _zero_grads(layers) + x.grad = None + dgrad_ref = _forward_backward(layers, x, recompute=False) + grads_ref = [l.weight.main_grad.clone() for l in recomputed] + + _zero_grads(layers) + x.grad = None + dgrad_rc = _forward_backward(layers, x, recompute=True) + grads_rc = [l.weight.main_grad.clone() for l in recomputed] + + # Both dgrad and wgrad flow through the recomputed weights, so both witness a clobber. + torch.testing.assert_close(dgrad_rc.float(), dgrad_ref.float(), rtol=1e-3, atol=1e-3) + for i, (g_rc, g_ref) in enumerate(zip(grads_rc, grads_ref)): + # A clobbered recompute uses a neighbour's weight, which moves the grad by O(grad), + # far outside this tolerance. + torch.testing.assert_close( + g_rc.float(), g_ref.float(), rtol=1e-3, atol=1e-3, msg=f"layer {i} wgrad mismatch" + ) + + +class TestGTPRecomputeChainBuffers: + def test_adjacent_nodes_use_distinct_buffers(self): + _requires_multi_gpu(4) + _run_distributed(_worker_adjacent_nodes_use_distinct_buffers, 4) + + +class TestGTPRecomputeCorrectness: + def test_recompute_matches_no_recompute(self): + _requires_multi_gpu(4) + _run_distributed(_worker_recompute_matches_no_recompute, 4) + + +# --------------------------------------------------------------------------- +# Grouped (routed-expert) recompute chains +# --------------------------------------------------------------------------- + +NUM_GEMMS = 2 +GROUPED_CHAIN = "GTP_remat_grouped_fc1_ungraphed" +# Strict subset: block 1 and 3 are gathered but never recomputed (see the worker docstring). +RECOMPUTED_BLOCKS = {0, 2} + + +def _worker_grouped_adjacent_nodes_use_distinct_buffers(rank, world_size, port): + """Same invariant on a grouped one-block-ahead chain, per expert. + + Two details make this a real test rather than a tautology: + * only a SUBSET of blocks is recomputed. The grouped chain's own _double_buffer_parity is + drawn in FORWARD order over every block, so recomputing all of them leaves the recompute + chain alternating by accident and the test passes even when unguarded. Skipping a block + makes two same-parity weights adjacent on the recompute chain -- the real collision. + * buffers are compared per EXPERT: grouped weights gather as a batch and the cache key + carries expert_idx, so expert k of block N collides with expert k of block N+1, not with + the anchor. + """ + torch.manual_seed(0) + gtp_module.reset_gtp_state() + + gtp_remat_group = dist.new_group(list(range(world_size))) + blocks = [ + _make_gtp_remat_grouped_linear( + NUM_GEMMS, HIDDEN, HIDDEN, gtp_remat_group, DTYPE, fuse_wgrad_accumulation=True + ) + for _ in range(NUM_LAYERS) + ] + # Production assigns these from the param name in _classify_param_chain; do it by hand so + # the weights land on the grouped one-block-ahead chain rather than the generic one. + for block in blocks: + for w in block.weight0.weight_list: + w.chain_id = GROUPED_CHAIN + w.main_grad = torch.zeros(w.shape, dtype=DTYPE, device="cuda") + + tokens = 8 * NUM_GEMMS + m_splits = [tokens // NUM_GEMMS] * NUM_GEMMS + x = torch.randn(tokens, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=True) + dist.broadcast(x, src=0) + + def fwd_bwd(): + h = x + for i, block in enumerate(blocks): + call = lambda inp, b=block: b(inp, m_splits=m_splits, is_first_microbatch=True) + if i in RECOMPUTED_BLOCKS: + checkpoint = CheckpointWithoutOutput() + y = checkpoint.checkpoint(call, h) + h = torch.nn.functional.gelu(y) + checkpoint.discard_output_and_register_recompute(h) + else: + h = torch.nn.functional.gelu(call(h)) + h.float().sum().backward() + + for _ in range(2): # first pass builds the chain, second uses it + fwd_bwd() + + chain = [blocks[i].weight0 for i in sorted(RECOMPUTED_BLOCKS)] + assert all( + a._recompute_initialized for a in chain + ), "grouped weights never gathered under recompute -- the chain was not built" + + failures = [] + for expert in range(NUM_GEMMS): + addrs = [_recompute_buffer_addr(a.weight_list[expert]) for a in chain] + if any(a is None for a in addrs): + failures.append(f"expert {expert}: unallocated buffer {addrs}") + continue + shared = [i for i in range(len(addrs) - 1) if addrs[i] == addrs[i + 1]] + if shared: + failures.append( + f"expert {expert}: chain nodes {shared} share a buffer with their successor " + f"({[hex(a) for a in addrs]})" + ) + assert not failures, "grouped recompute chain collides:\n " + "\n ".join(failures) + + +class TestGroupedGTPRecomputeChainBuffers: + def test_grouped_adjacent_nodes_use_distinct_buffers(self): + _requires_multi_gpu(4) + _run_distributed(_worker_grouped_adjacent_nodes_use_distinct_buffers, 4) diff --git a/tests/unit_tests/models/test_gpt_model.py b/tests/unit_tests/models/test_gpt_model.py index d2cb12841c4..64b8deeb9c6 100644 --- a/tests/unit_tests/models/test_gpt_model.py +++ b/tests/unit_tests/models/test_gpt_model.py @@ -144,6 +144,7 @@ def test_output_processor_forward(self): ).cuda() context = {"selected_token_positions": torch.tensor([0, 2], device="cuda")} + created = {} seen = {} def output_processor(**kwargs): @@ -158,7 +159,14 @@ def output_processor(**kwargs): dim=-1, index=kwargs["labels"].unsqueeze(-1) ) token_logprobs = token_logprobs.squeeze(-1) - return token_logprobs.index_select(1, kwargs["context"]["selected_token_positions"]) + result = { + "payload": token_logprobs.index_select( + 1, kwargs["context"]["selected_token_positions"] + ), + "tag": "structured", + } + created["result"] = result + return result with torch.no_grad(): logits = self.gpt_model.forward( @@ -176,10 +184,14 @@ def output_processor(**kwargs): output_processor_context=context, ) - assert torch.allclose(output, expected) + assert output is created["result"] + assert isinstance(output, dict) + assert torch.allclose(output["payload"], expected) + assert output["tag"] == "structured" assert seen["context"] is context assert seen["output_layer"] is self.gpt_model.output_layer assert seen["output_weight"] is None + assert seen["output_layer"].weight is not None assert seen["labels"] is labels assert seen["runtime_gather_output"] is None assert seen["config"] is config diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index c37d5ec00e1..72b65689ff5 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -108,16 +108,37 @@ def make_token_rollout(trajectory, logprobs, generation_mask=None, reward=1.0, p ) -class DummyLangModule: +class DummyConfigModule(torch.nn.Module): def __init__(self, config): + super().__init__() + self.config = config + + +class DummyLogprobsModel(torch.nn.Module): + def __init__(self, config, layer_config): + super().__init__() + self.config = config + self.layer = DummyConfigModule(layer_config) + self.pg_collection = SimpleNamespace(pp=object()) + self.config_values_during_forward = None + + def forward(self, tokens, position_ids, attention_mask, **kwargs): + del position_ids, attention_mask, kwargs + self.config_values_during_forward = ( + self.config.flash_decode, + self.layer.config.flash_decode, + ) + return torch.ones((tokens.shape[0], tokens.shape[1], VOCAB)) + + +class DummyLangModule(torch.nn.Module): + def __init__(self, config): + super().__init__() self.config = config self.rotary_pos_emb = None self.eval = MagicMock() self.train = MagicMock() - def modules(self): - return iter(()) - class DummyMoELayer: def __init__(self, use_partial_cudagraphs): @@ -346,17 +367,33 @@ def _toggle(lang_module, set_to): return MagicMock(side_effect=_toggle) - def test_megatron_rl_inference_mode_restores_training_cuda_graph_state(self, monkeypatch): + @pytest.mark.parametrize( + "share_config", + [pytest.param(True, id="shared-config"), pytest.param(False, id="distinct-config")], + ) + @pytest.mark.parametrize("num_experts", [None, 8], ids=["dense", "moe"]) + def test_megatron_rl_inference_mode_restores_training_cuda_graph_state( + self, monkeypatch, share_config, num_experts + ): config = SimpleNamespace( cuda_graph_impl="none", cuda_graph_modules=[CudaGraphModule.attn], inference_cuda_graph_scope=InferenceCudaGraphScope.none, ) - lang_module = DummyLangModule(config) + layer_config = ( + config + if share_config + else SimpleNamespace( + cuda_graph_impl="none", + cuda_graph_modules=[CudaGraphModule.attn], + inference_cuda_graph_scope=InferenceCudaGraphScope.none, + ) + ) + lang_module = DummyLangModule(layer_config) model = [SimpleNamespace(config=config, module=lang_module)] args = SimpleNamespace( rl_training_cuda_graphs=False, - num_experts=None, + num_experts=num_experts, curr_iteration=11, cuda_graph_impl="local", cuda_graph_modules=[CudaGraphModule.attn], @@ -368,20 +405,61 @@ def test_megatron_rl_inference_mode_restores_training_cuda_graph_state(self, mon with rl_utils.megatron_rl_inference_mode(model, MagicMock(), "local", False) as result: assert result is interface - assert config.cuda_graph_impl == "local" - assert config.cuda_graph_modules == [] - assert config.inference_cuda_graph_scope == InferenceCudaGraphScope.block + for current_config in (config, layer_config): + assert current_config.cuda_graph_impl == "local" + assert current_config.cuda_graph_modules == [] + assert current_config.inference_cuda_graph_scope == InferenceCudaGraphScope.block assert toggle_cuda_graphs.call_args_list == [ call(lang_module, "local"), call(lang_module, "none"), ] - assert config.cuda_graph_impl == "local" - assert config.cuda_graph_modules == [CudaGraphModule.attn] - assert config.inference_cuda_graph_scope == InferenceCudaGraphScope.block + expected_modules = ( + [ + CudaGraphModule.mamba, + CudaGraphModule.attn, + CudaGraphModule.moe_router, + CudaGraphModule.moe_preprocess, + ] + if num_experts is not None + else [CudaGraphModule.attn] + ) + for current_config in (config, layer_config): + assert current_config.cuda_graph_impl == "local" + assert current_config.cuda_graph_modules == expected_modules + assert current_config.inference_cuda_graph_scope == InferenceCudaGraphScope.block lang_module.eval.assert_called_once() lang_module.train.assert_called_once() + @pytest.mark.parametrize( + "share_config", + [pytest.param(True, id="shared-config"), pytest.param(False, id="distinct-config")], + ) + def test_get_logprobs_updates_all_model_configs(self, monkeypatch, share_config): + config = SimpleNamespace(flash_decode=True) + layer_config = config if share_config else SimpleNamespace(flash_decode=True) + model = DummyLogprobsModel(config, layer_config) + monkeypatch.setattr(rl_utils, "get_args", lambda: SimpleNamespace(fp16=False, bf16=False)) + monkeypatch.setattr( + rl_utils, "get_nvtx_range", lambda: (lambda *args, **kwargs: nullcontext()) + ) + monkeypatch.setattr( + rl_utils, "get_attr_wrapped_model", lambda model, name: getattr(model, name) + ) + monkeypatch.setattr(rl_utils, "is_pp_last_stage", lambda _group: False) + + output = rl_utils.get_logprobs( + model, + torch.ones((1, 2), dtype=torch.long), + position_ids=None, + packed_seq_params=object(), + ) + + assert output.shape == (1, 2, VOCAB) + assert model.config_values_during_forward == (False, False) + assert config.flash_decode is True + assert layer_config.flash_decode is True + @pytest.mark.parametrize( "initialize_model_parallel", [ diff --git a/tests/unit_tests/ssm/test_gdp_dynamic_inference.py b/tests/unit_tests/ssm/test_gdp_dynamic_inference.py new file mode 100644 index 00000000000..cbb4fbc3821 --- /dev/null +++ b/tests/unit_tests/ssm/test_gdp_dynamic_inference.py @@ -0,0 +1,629 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GDP (Gated Delta Product) inference equivalence tests. + +`GatedDeltaProductMixer` supports two inference paths: static batching +(`StaticInferenceContext`) and dynamic batching (`DynamicInferenceContext`). +These tests assert that both are numerically equivalent to a plain +full-sequence forward, and therefore to each other. + +The reference is a single full-sequence `model.forward` (the +`chunk_gated_delta_product` path). Dynamic prefill runs the same chunk kernel +over a packed var-len layout, and static-batching prefill runs the same chunk +kernel with the recurrent cache seeded; both must reproduce the reference's +last-token logits. + +The single-forward equivalence tests run at TP=1 (they compare raw logits that +are sequence-sharded under sequence-parallel, and the static path does not +support SP). The end-to-end engine tests sweep TP (`_TP_SIZES`) with SP enabled +at TP>1, covering dynamic inference under tensor + sequence parallelism; TP>1 +variants skip when the world has too few GPUs. +""" + +from __future__ import annotations + +import random +import types +from typing import Dict, List, Optional, Sequence, Tuple + +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig +from megatron.core.inference.contexts import StaticInferenceContext +from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext +from megatron.core.inference.engines import DynamicInferenceEngine +from megatron.core.inference.inference_request import DynamicInferenceRequest, Status +from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import ( + GPTInferenceWrapper, +) +from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) +from megatron.core.inference.utils import InferenceMode +from megatron.core.models.hybrid.hybrid_layer_specs import gated_delta_product_inference_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer +from megatron.core.ssm.packed_seq_helpers import check_fla_sequence_packing_support +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.cuda_graphs import delete_cuda_graphs +from megatron.core.utils import is_fa_min_version +from tests.unit_tests.test_utilities import Utils, clear_nvte_env_vars + +try: + import einops # noqa: F401 + import fla # noqa: F401 + import mamba_ssm # noqa: F401 + + HAVE_GDP_DEPS = True +except ImportError: + HAVE_GDP_DEPS = False + +# GDP dynamic inference relies on the same packed-sequence conv1d kernel as the +# training/prefill path (`causal_conv1d_fn(seq_idx=...)`, added in 1.4.0). +_PACKING_OK, _PACKING_REASON = check_fla_sequence_packing_support() + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif(not HAVE_GDP_DEPS, reason="GDP requires fla, mamba_ssm, and einops"), + pytest.mark.skipif(not _PACKING_OK, reason=_PACKING_REASON or "packed-seq support missing"), +] + + +# A short single-chunk prompt is enough to exercise the packed-varlen dynamic +# path; the sizes are kept small to keep the test fast. +_VOCAB_SIZE = 128 +_MAX_SEQ_LEN = 512 +_PROMPT_LEN = 64 + +# bf16 chunk-vs-chunk tolerance. Full-forward and prefill run the same +# `chunk_gated_delta_product` kernel, so they differ only by the packed var-len +# layout and floating-point accumulation order. +_ATOL = 5e-2 +_RTOL = 5e-2 + +# Looser tolerance for the decode-step check: it compares the recurrent decode +# kernel against a full-sequence chunk-kernel recompute (different kernels), so +# it drifts more than the chunk-vs-chunk prefill comparison. Still far tighter +# than the O(1)+ deviations a genuinely broken decode/state-handoff would show. +_DECODE_ATOL = 1e-1 +_DECODE_RTOL = 1e-1 + +# `DynamicInferenceContext` requires at least one attention layer, so the model +# pattern is GDP mixer + attention + MLP. Dynamic batching needs a recent +# flash-attention. +_LAYER_PATTERN = "M*-" +_NUM_LAYERS = len(_LAYER_PATTERN) +requires_dynamic_batching = pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need flash-attn >= 2.7.3 for dynamic batching" +) + +# Tensor-parallel sizes swept by the end-to-end engine tests. TP>1 requires +# sequence-parallel (the inference-optimized linears assert it), which GDP's +# dynamic path supports; static inference does not, so the single-forward +# equivalence tests above stay TP=1. +_TP_SIZES = [1, 2] + + +def _make_config(tp: int = 1) -> TransformerConfig: + """A small but shape-valid GDP config, sharded across `tp` tensor-parallel ranks. + + The in_proj output width (`zVKQba`) is column-parallel, so each rank sees + `proj_dim / tp` channels. The packed-prefill conv slices a channels-last view + out of that and `causal_conv1d_fn` requires its stride (the per-rank width) to + be a multiple of 8. With mamba_num_heads=16 the full width is + `(1+M)*d_inner + (M+1)*ngroups*d_state + (M+1)*nheads = 3*256 + 3*64 + 3*16 + = 1008`, so per-rank widths are 1008 (tp=1) and 504 (tp=2), both aligned. + Production configs satisfy this by having much larger, aligned dimensions. + """ + return TransformerConfig( + num_layers=_NUM_LAYERS, + hidden_size=64, + num_attention_heads=4, + num_query_groups=4, + ffn_hidden_size=128, + normalization="RMSNorm", + bf16=True, + params_dtype=torch.bfloat16, + mamba_num_heads=16, + mamba_head_dim=16, + mamba_num_groups=4, + mamba_state_dim=16, + gdp_num_householder=2, + is_hybrid_model=True, # needed for correct out_proj init + tensor_model_parallel_size=tp, + sequence_parallel=False, + context_parallel_size=1, + ) + + +def _build_model(tp: int = 1) -> HybridModel: + """Build a small GDP hybrid model (mixer + attention + MLP), eval on CUDA.""" + model_parallel_cuda_manual_seed(123) + model = HybridModel( + config=_make_config(tp), + hybrid_stack_spec=gated_delta_product_inference_stack_spec, + vocab_size=_VOCAB_SIZE, + max_sequence_length=_MAX_SEQ_LEN, + hybrid_layer_pattern=_LAYER_PATTERN, + ) + return model.cuda().eval() + + +@requires_dynamic_batching +class TestGDPDynamicInference: + """Static/dynamic GDP inference equivalence against a full-sequence forward. + + These compare raw `model.forward` logits, which are sequence-sharded under + sequence-parallel; combined with the static path not supporting SP, they run + at TP=1 only. TP>1 dynamic inference is covered end-to-end by the engine + tests below, which handle SP through the inference wrapper. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + self.model = _build_model(tp=1) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _input_ids(self) -> torch.Tensor: + """A deterministic single-request prompt: shape [1, _PROMPT_LEN].""" + return torch.arange(_PROMPT_LEN, device="cuda", dtype=torch.long).unsqueeze(0) + + @torch.inference_mode() + def _full_forward_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Reference: plain full-sequence forward -> last-token logits [1, V].""" + # No inference_context: GDP runs the training / chunk_gated_delta_product + # path. This is the ground truth both inference modes must reproduce. + InferenceMode.unset_active() + position_ids = torch.arange(input_ids.shape[1], device="cuda").unsqueeze(0) + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + runtime_gather_output=True, + ) + return logits[:, -1, :].float() + + def _build_dynamic_context(self) -> DynamicInferenceContext: + mamba_config = MambaInferenceStateConfig.from_model(self.model) + assert mamba_config is not None, "GDP hybrid model should expose Mamba inference state" + return DynamicInferenceContext( + model_config=self.model.config, + inference_config=InferenceConfig( + max_sequence_length=_MAX_SEQ_LEN, + buffer_size_gb=1.0, + block_size_tokens=256, + # Materialize all tokens so we can read the prompt's last-token + # logits directly (static batching always uses last-token only). + materialize_only_last_token_logits=False, + mamba_inference_state_config=mamba_config, + num_cuda_graphs=0, + use_cuda_graphs_for_non_decode_steps=False, + max_requests=4, + max_tokens=128, + ), + ) + + @torch.inference_mode() + def _dynamic_prefill_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Dynamic-batching prefill -> last-token logits [1, V].""" + ctx = self._build_dynamic_context() + request = DynamicInferenceRequest( + request_id=0, + prompt_tokens=input_ids.cpu().squeeze(0), + sampling_params=SamplingParams(num_tokens_to_generate=1, termination_id=-1), + ) + ctx.add_request(request) + ctx.initialize_attention_state() + with InferenceMode.active(): + logits = self.model( + input_ids=input_ids, + position_ids=None, + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + # materialize_only_last_token_logits=False -> [1, prompt_len, V]. + return logits[:, -1, :].float() + + @torch.inference_mode() + def _static_prefill_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Static-batching prefill -> last-token logits [1, V].""" + ctx = StaticInferenceContext(max_batch_size=1, max_sequence_length=_MAX_SEQ_LEN) + ctx.sequence_len_offset = 0 + position_ids = torch.arange(input_ids.shape[1], device="cuda").unsqueeze(0) + with InferenceMode.active(): + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + # StaticInferenceContext forces materialize_only_last_token_logits=True, + # so the sequence dimension is already collapsed to the last token. + assert logits.shape[1] == 1 + return logits[:, 0, :].float() + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + def test_constructor(self): + """The GDP stack spec wires a GatedDeltaProductMixer into the mamba layer.""" + assert isinstance(self.model, HybridModel) + mixers = [ + layer.mixer + for layer in self.model.decoder.layers + if hasattr(layer, "mixer") and layer.mixer is not None + ] + assert len(mixers) == 1, f"pattern {_LAYER_PATTERN!r} should yield exactly one mixer layer" + assert isinstance(mixers[0], GatedDeltaProductMixer) + + def test_full_forward_shape(self): + """Sanity check: plain forward returns [batch, seq, vocab].""" + input_ids = self._input_ids() + InferenceMode.unset_active() + position_ids = torch.arange(_PROMPT_LEN, device="cuda").unsqueeze(0) + with torch.inference_mode(): + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + runtime_gather_output=True, + ) + assert logits.shape == (1, _PROMPT_LEN, _VOCAB_SIZE) + + def test_dynamic_prefill_matches_full_forward(self): + """Dynamic-batching prefill reproduces the full-sequence forward.""" + input_ids = self._input_ids() + reference = self._full_forward_last_logits(input_ids) + dynamic = self._dynamic_prefill_last_logits(input_ids) + torch.testing.assert_close(dynamic, reference, atol=_ATOL, rtol=_RTOL) + + def test_static_prefill_matches_full_forward(self): + """Static-batching prefill reproduces the full-sequence forward.""" + input_ids = self._input_ids() + reference = self._full_forward_last_logits(input_ids) + static = self._static_prefill_last_logits(input_ids) + torch.testing.assert_close(static, reference, atol=_ATOL, rtol=_RTOL) + + def test_static_and_dynamic_prefill_agree(self): + """Static and dynamic inference produce equivalent logits. + + This is the central invariant: the two batching strategies must agree. + Anchoring each to the full-sequence forward (above) guarantees this + transitively, but assert it directly as well so a regression in either + path that happens to drift in the same direction is still caught. + """ + input_ids = self._input_ids() + static = self._static_prefill_last_logits(input_ids) + dynamic = self._dynamic_prefill_last_logits(input_ids) + torch.testing.assert_close(static, dynamic, atol=_ATOL, rtol=_RTOL) + + @torch.inference_mode() + def test_decode_step_matches_recompute(self): + """One decode step matches a full-sequence recompute (decode-path check). + + This validates the recurrent decode kernel and the prefill->decode + conv/SSM state handoff independently of any golden snapshot: after + prefilling the prompt, decoding one more token must produce the same + next-token logits as a plain forward over prompt+token. Compared at the + logit level with tolerance, so it is robust to the bf16 numerics that + make exact greedy token equality across the recurrent/chunk kernels + fragile. Uses `StaticInferenceContext`, whose decode calls the same + `ssm_decode` recurrent kernel as the dynamic engine. + """ + prompt = self._input_ids() # [1, P] + + # Ground truth: the next token from the prompt, and the full-recompute + # distribution for the token after it (chunk kernel over prompt+token). + next_token = int(self._full_forward_last_logits(prompt).argmax(dim=-1).item()) + extended = torch.cat( + [prompt, torch.tensor([[next_token]], dtype=torch.int64, device="cuda")], dim=1 + ) + recompute = self._full_forward_last_logits(extended) # [1, V] + + # Incremental path: prefill the prompt, then a single decode step. + ctx = StaticInferenceContext(max_batch_size=1, max_sequence_length=_MAX_SEQ_LEN) + prompt_length = prompt.shape[1] + with InferenceMode.active(): + ctx.sequence_len_offset = 0 + self.model( + input_ids=prompt, + position_ids=torch.arange(prompt_length, device="cuda").unsqueeze(0), + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + ctx.sequence_len_offset = prompt_length + decode_logits = self.model( + input_ids=torch.tensor([[next_token]], dtype=torch.int64, device="cuda"), + position_ids=torch.tensor([[prompt_length]], dtype=torch.int64, device="cuda"), + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + decode_last = decode_logits[:, -1, :].float() + + torch.testing.assert_close(decode_last, recompute, atol=_DECODE_ATOL, rtol=_DECODE_RTOL) + + +# ====================================================================== +# End-to-end engine tests. +# +# The tests above exercise a single forward pass. These drive the full +# `DynamicInferenceEngine` (add requests -> schedule -> prefill -> decode -> +# finish) so that GDP is validated through the same runtime path production +# inference uses: the text-generation controller, the inference-wrapped model, +# the KV/Mamba-state cache, and the request scheduler. +# +# Decoding is greedy (`top_k=1`) so outputs are deterministic. Decode +# correctness (the recurrent kernel `fused_recurrent_gated_delta_rule` plus the +# slot-indexed conv/SSM cache, distinct from the prefill chunk kernel) is +# validated by a byte-for-byte match against committed golden token ids, +# mirroring the Mamba2 `test_dynamic_engine.py::test_simple` style. The golden +# constant is captured from a reference GPU run (see `_GOLDEN_*` below); until it +# is populated the test self-captures and skips with the observed ids. +# ====================================================================== + +# Fixed prompts for the golden-token test. Deterministic (not random) so the +# committed golden ids below are reproducible across machines. Varying lengths +# exercise the scheduler's mixed-length prefill batching. +_GOLDEN_PROMPTS: List[List[int]] = [ + [3, 14, 15, 92, 65, 35, 89, 79], + [2, 71, 82, 81, 8], + [11, 22, 33, 44, 55, 66], + [7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7], +] +_GOLDEN_NUM_TOKENS_TO_GENERATE = 12 + +# Golden generated-token ids per TP size, one list per prompt in `_GOLDEN_PROMPTS`, +# captured from a reference GPU run. TP shards the weights differently, so each TP +# size has its own goldens. Environment-sensitive (FLA / causal_conv1d kernel +# build, GPU arch); re-capture if the kernels or config change. While an entry is +# None, the test self-captures: it prints the observed ids and skips instead of +# failing. Paste them in (see the skip message) to turn it into a hard assertion. +_GOLDEN_GENERATED_TOKENS: Dict[int, Optional[List[List[int]]]] = { + 1: [ + [32, 126, 35, 125, 52, 116, 55, 38, 39, 4, 53, 100], + [88, 10, 105, 95, 105, 44, 2, 100, 127, 59, 23, 18], + [29, 9, 61, 2, 100, 2, 69, 75, 36, 80, 103, 26], + [2, 55, 4, 108, 116, 120, 24, 113, 100, 48, 111, 22], + ], + 2: [ + [12, 51, 63, 22, 2, 40, 45, 30, 55, 10, 31, 29], + [38, 54, 29, 17, 33, 13, 10, 45, 1, 22, 21, 37], + [24, 55, 38, 55, 35, 61, 26, 25, 31, 20, 62, 56], + [53, 26, 31, 61, 16, 19, 42, 41, 49, 18, 53, 26], + ], +} + + +def _make_engine_config(tp: int = 1) -> TransformerConfig: + """GDP config for the engine tests: same shape as `_make_config`, plus the + deterministic inference sampling knobs greedy decoding needs. TP>1 turns on + sequence-parallel, which the inference-optimized linears require.""" + config = _make_config(tp) + config.sequence_parallel = tp > 1 + config.inference_rng_tracker = True + config.inference_sampling_seed = 123 + return config + + +@pytest.mark.internal +@requires_dynamic_batching +@pytest.mark.skipif(not HAVE_GDP_DEPS, reason="GDP requires fla, mamba_ssm, and einops") +@pytest.mark.skipif(not _PACKING_OK, reason=_PACKING_REASON or "packed-seq support missing") +class TestGDPDynamicInferenceEngine: + """End-to-end GDP decoding through `DynamicInferenceEngine`.""" + + SEED = 123 + VOCAB_SIZE = _VOCAB_SIZE + + def teardown_method(self, method): + delete_cuda_graphs() + Utils.destroy_model_parallel() + + # ------------------------------------------------------------------ + # Harness + # ------------------------------------------------------------------ + + def _build_engine( + self, + *, + num_tokens_to_generate: int, + tp: int = 1, + num_requests: Optional[int] = None, + prompt_length: Optional[int] = None, + prompts: Optional[Sequence[Sequence[int]]] = None, + ) -> Tuple[DynamicInferenceEngine, List[DynamicInferenceRequest]]: + """Build a greedy GDP engine plus its requests at TP=`tp`. + + Either pass explicit `prompts` (deterministic token lists) or + `num_requests` + `prompt_length` (random prompts of a fixed length). + Skips if the world is too small for the requested TP size. + """ + if Utils.world_size < tp: + pytest.skip(f"TP={tp} requires at least {tp} GPUs") + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp, pipeline_model_parallel_size=1 + ) + clear_nvte_env_vars() + random.seed(self.SEED) + torch.manual_seed(self.SEED) + model_parallel_cuda_manual_seed( + seed=self.SEED, inference_rng_tracker=True, force_reset_rng=True + ) + + if prompts is not None: + prompt_tensors = [torch.tensor(p, dtype=torch.int64, device="cuda") for p in prompts] + else: + assert num_requests is not None and prompt_length is not None + prompt_tensors = [ + torch.randint( + 0, self.VOCAB_SIZE - 1, (prompt_length,), dtype=torch.int64, device="cuda" + ) + for _ in range(num_requests) + ] + + max_prompt_length = max(int(p.numel()) for p in prompt_tensors) + max_sequence_length = max_prompt_length + num_tokens_to_generate + config = _make_engine_config(tp) + model = HybridModel( + config=config, + hybrid_stack_spec=gated_delta_product_inference_stack_spec, + vocab_size=self.VOCAB_SIZE, + max_sequence_length=max_sequence_length, + parallel_output=True, + hybrid_layer_pattern=_LAYER_PATTERN, + pre_process=parallel_state.is_pipeline_first_stage(), + post_process=parallel_state.is_pipeline_last_stage(), + ).cuda() + for param in model.parameters(): + param.data = param.data.to(config.params_dtype) + model.eval() + + context = DynamicInferenceContext( + model_config=config, + inference_config=InferenceConfig( + max_sequence_length=max_sequence_length, + buffer_size_gb=0.1, + block_size_tokens=256, + materialize_only_last_token_logits=True, + mamba_inference_state_config=MambaInferenceStateConfig.from_model(model), + num_cuda_graphs=None, + use_cuda_graphs_for_non_decode_steps=False, + max_requests=32, + max_tokens=1024, + ), + ) + + wrapped_model = GPTInferenceWrapper(model, context) + wrapped_model.model_is_pipeline_parallel = not ( + parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage() + ) + controller = TextGenerationController( + inference_wrapped_model=wrapped_model, + tokenizer=types.SimpleNamespace( + vocab_size=self.VOCAB_SIZE, detokenize=lambda tokens: "tokenized_prompt" + ), + ) + delete_cuda_graphs() + engine = DynamicInferenceEngine(controller, context) + + requests = [ + DynamicInferenceRequest( + request_id=request_id, + prompt_tokens=prompt_tokens, + sampling_params=SamplingParams( + num_tokens_to_generate=num_tokens_to_generate, + termination_id=-1, # never terminate early -> fixed output length + top_k=1, # greedy -> deterministic + ), + ) + for request_id, prompt_tokens in enumerate(prompt_tensors) + ] + return engine, requests + + @staticmethod + @torch.inference_mode() + def _run_to_completion( + engine: DynamicInferenceEngine, requests: List[DynamicInferenceRequest] + ) -> Dict[int, DynamicInferenceRequest]: + """Add every request, step until the engine drains, return finished requests by id.""" + for request in requests: + engine._add_request(request) + + finished: Dict[int, DynamicInferenceRequest] = {} + # Bound the loop so a scheduling regression fails loudly instead of hanging. + for _ in range(1000): + result = engine.step_modern() + for record in result["finished_request_records"]: + merged = record.merge() + finished[merged.request_id] = merged + if not engine.has_unfinished_requests(): + break + assert not engine.has_unfinished_requests(), "engine did not drain within step budget" + return finished + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_engine_runs_to_completion(self, tp): + """Every request completes and yields exactly the requested token count.""" + num_tokens_to_generate = 8 + engine, requests = self._build_engine( + tp=tp, num_requests=4, prompt_length=8, num_tokens_to_generate=num_tokens_to_generate + ) + finished = self._run_to_completion(engine, requests) + + assert len(finished) == len(requests) + for request in requests: + merged = finished[request.request_id] + assert merged.status == Status.COMPLETED + # termination_id=-1 disables early stop, so the length is exact. + assert len(merged.generated_tokens) == num_tokens_to_generate + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_engine_greedy_matches_golden(self, tp): + """Greedy decode reproduces committed golden token ids (Mamba2-style). + + Deterministic fixed prompts + greedy sampling make the output a stable + fingerprint of the GDP prefill+decode path. Until the TP entry in + `_GOLDEN_GENERATED_TOKENS` is captured from a reference GPU run, the test + prints the observed ids and skips instead of failing. + """ + engine, requests = self._build_engine( + tp=tp, prompts=_GOLDEN_PROMPTS, num_tokens_to_generate=_GOLDEN_NUM_TOKENS_TO_GENERATE + ) + finished = self._run_to_completion(engine, requests) + observed = [finished[r.request_id].generated_tokens for r in requests] + + golden = _GOLDEN_GENERATED_TOKENS.get(tp) + if golden is None: + pytest.skip( + f"golden tokens for TP={tp} not captured yet; paste the following into " + f"_GOLDEN_GENERATED_TOKENS[{tp}]:\n{observed!r}" + ) + + assert observed == golden, ( + f"generated tokens != golden (TP={tp}):\n golden = {golden}\n" + f" observed = {observed}" + ) + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_generate_over_multiple_prompts(self, tp): + """`engine.generate` drives several prompts through to completion at once.""" + engine, requests = self._build_engine( + tp=tp, num_requests=4, prompt_length=8, num_tokens_to_generate=4 + ) + + prompts = [f"prompt{i}" for i in range(len(requests))] + + def mock_tokenize_prompt(tokenizer, prompt, add_BOS=False): + prompt_num = int(prompt[-1]) + return [10 + i for i in range(prompt_num + 2)] + + engine.controller.tokenize_prompt = mock_tokenize_prompt + + finished_records = engine.generate(prompts, requests[0].sampling_params) + finished = [record.merge() for record in finished_records] + + assert len(finished) == len(prompts) + # generate() returns finished requests in request-id order. + assert [r.request_id for r in finished] == sorted(r.request_id for r in finished) + for request in finished: + assert request.status == Status.COMPLETED + assert len(request.generated_tokens) > 0 diff --git a/tests/unit_tests/tensor_parallel/test_layers.py b/tests/unit_tests/tensor_parallel/test_layers.py index dbc27f502c6..cf1e8185017 100644 --- a/tests/unit_tests/tensor_parallel/test_layers.py +++ b/tests/unit_tests/tensor_parallel/test_layers.py @@ -2,11 +2,33 @@ import pytest import torch -from megatron.core.tensor_parallel.layers import linear_with_frozen_weight +from megatron.core.extensions.transformer_engine import te_general_gemm +from megatron.core.tensor_parallel.layers import ( + linear_with_frozen_weight, + linear_with_grad_accumulation_and_async_allreduce, +) from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region from tests.unit_tests.test_utilities import Utils +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_linear_default_output_dtype_preserves_input_dtype(dtype): + Utils.initialize_model_parallel(1, 1) + + try: + input_data = torch.randn(4, 3, 16, device="cuda", dtype=dtype) + weight = torch.randn(32, 16, device="cuda", dtype=dtype) + output = linear_with_grad_accumulation_and_async_allreduce( + input_data, weight, None, False, False, False, tp_group=None, output_dtype=None + ) + reference = torch.nn.functional.linear(input_data, weight) + + assert output.dtype == input_data.dtype + torch.testing.assert_close(output, reference) + finally: + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("tensor_parallel,allreduce_dgrad", [(1, False), (8, True)]) def test_LinearWithFrozenWeight(tensor_parallel, allreduce_dgrad): Utils.initialize_model_parallel(tensor_parallel, 1) @@ -70,3 +92,111 @@ def test_LinearWithFrozenWeight_3d_input_matches_torch_linear(): assert torch.allclose(input_data.grad, expected_input.grad) Utils.destroy_model_parallel() + + +@pytest.mark.skipif( + te_general_gemm is None, reason="Transformer Engine general_gemm is not available" +) +def test_linear_with_grad_accumulation_supports_fp32_output_and_bf16_backward(): + Utils.initialize_model_parallel(1, 1) + + input_data = torch.randn(4, 3, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(32, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + reference_input = input_data.detach().clone().requires_grad_(True) + reference_weight = weight.detach().clone().requires_grad_(True) + + output = linear_with_grad_accumulation_and_async_allreduce( + input_data, weight, None, False, False, False, tp_group=None, output_dtype=torch.float32 + ) + output.sum().backward() + + reference_output = torch.nn.functional.linear(reference_input, reference_weight) + reference_output.sum().backward() + fp32_reference_output = torch.nn.functional.linear( + input_data.detach().float(), weight.detach().float() + ) + + assert output.dtype == torch.float32 + assert input_data.grad.dtype == torch.bfloat16 + assert weight.grad.dtype == torch.bfloat16 + assert torch.allclose(output, fp32_reference_output, atol=1e-4, rtol=1e-4) + assert torch.allclose(input_data.grad, reference_input.grad) + assert torch.allclose(weight.grad, reference_weight.grad) + + Utils.destroy_model_parallel() + + +@pytest.mark.skipif( + te_general_gemm is None, reason="Transformer Engine general_gemm is not available" +) +def test_linear_fp32_output_is_bitwise_exact_for_integer_bf16_operands(): + Utils.initialize_model_parallel(1, 1) + + generator = torch.Generator(device="cuda").manual_seed(1234) + input_data = torch.randint( + -8, 9, (4, 3, 512), device="cuda", dtype=torch.int32, generator=generator + ).to(torch.bfloat16) + weight = torch.randint( + -8, 9, (128, 512), device="cuda", dtype=torch.int32, generator=generator + ).to(torch.bfloat16) + + output = linear_with_grad_accumulation_and_async_allreduce( + input_data, weight, None, False, False, False, tp_group=None, output_dtype=torch.float32 + ) + reference = torch.nn.functional.linear(input_data.float(), weight.float()) + + # K * 8^2 = 32,768, so every possible integer product and partial sum is + # exactly representable in FP32. Compare raw words instead of using a tolerance. + assert output.dtype == torch.float32 + assert torch.equal(output.contiguous().view(torch.int32), reference.view(torch.int32)) + + Utils.destroy_model_parallel() + + +@pytest.mark.skipif( + te_general_gemm is None, reason="Transformer Engine general_gemm is not available" +) +def test_linear_fp32_output_matches_plain_te_general_gemm(): + from transformer_engine.pytorch.cpp_extensions import general_gemm + + try: + from transformer_engine.pytorch.module.base import get_workspace + except ImportError: + get_workspace = None + + Utils.initialize_model_parallel(1, 1) + + input_data = torch.randn(4, 3, 64, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(96, 64, device="cuda", dtype=torch.bfloat16) + wrapped_output = linear_with_grad_accumulation_and_async_allreduce( + input_data, weight, None, False, False, False, tp_group=None, output_dtype=torch.float32 + ) + + kwargs = { + "out_dtype": torch.float32, + "quantization_params": None, + "gelu": None, + "gelu_in": None, + "accumulate": False, + "layout": "TN", + "out": None, + "bias": None, + "use_split_accumulator": False, + "grad": False, + "ub": None, + "ub_type": None, + "extra_output": None, + "bulk_overlap": False, + } + if get_workspace is not None: + kwargs["workspace"] = get_workspace() + plain_te_output = general_gemm(weight, input_data.reshape(-1, 64), **kwargs)[0] + plain_te_output = plain_te_output.reshape_as(wrapped_output) + + assert wrapped_output.dtype == torch.float32 + assert torch.equal( + wrapped_output.contiguous().view(torch.int32), + plain_te_output.contiguous().view(torch.int32), + ) + + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/test_argument_utils.py b/tests/unit_tests/test_argument_utils.py index 7c0b30d3d56..57a327cddfd 100644 --- a/tests/unit_tests/test_argument_utils.py +++ b/tests/unit_tests/test_argument_utils.py @@ -678,6 +678,29 @@ def test_transformer_callback_fields_are_not_registered_as_cli_args(self): assert not hasattr(args, field_name) +class TestMegatronMixedPrecisionArguments: + """Test language-model logit dtype CLI choices.""" + + @staticmethod + def _parser() -> ArgumentParser: + from megatron.training.arguments import _add_mixed_precision_args + + return _add_mixed_precision_args(ArgumentParser(exit_on_error=False)) + + def test_logit_dtype_defaults_to_input_dtype(self): + args = self._parser().parse_args([]) + assert args.logit_dtype is None + + @pytest.mark.parametrize("dtype", ["bf16", "fp32"]) + def test_logit_dtype_accepts_supported_choices(self, dtype): + args = self._parser().parse_args(["--output-logit-dtype", dtype]) + assert args.logit_dtype == dtype + + def test_logit_dtype_rejects_fp16(self): + with pytest.raises(ArgumentError, match="invalid choice"): + self._parser().parse_args(["--output-logit-dtype", "fp16"]) + + # --------------------------------------------------------------------------- # Tests for pretrain_cfg_container_from_args # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/training/models/test_gpt_builder.py b/tests/unit_tests/training/models/test_gpt_builder.py index 2263525e030..5ae7ab4f418 100644 --- a/tests/unit_tests/training/models/test_gpt_builder.py +++ b/tests/unit_tests/training/models/test_gpt_builder.py @@ -261,6 +261,7 @@ def test_default_values(self): assert config.should_pad_vocab is False assert config.seq_length == 1024 assert config.fp16_lm_cross_entropy is False + assert config.logit_dtype is None assert config.parallel_output is True assert config.share_embeddings_and_output_weights is False assert config.position_embedding_type == "learned_absolute" @@ -279,6 +280,7 @@ def test_custom_initialization(self): transformer=_make_transformer(), seq_length=4096, fp16_lm_cross_entropy=True, + logit_dtype=torch.float32, parallel_output=False, share_embeddings_and_output_weights=True, position_embedding_type="rope", @@ -287,6 +289,7 @@ def test_custom_initialization(self): ) assert config.seq_length == 4096 assert config.fp16_lm_cross_entropy is True + assert config.logit_dtype == torch.float32 assert config.parallel_output is False assert config.share_embeddings_and_output_weights is True assert config.position_embedding_type == "rope" @@ -678,6 +681,7 @@ def test_config_params_passed_to_mcore(self, mock_model, *_): vocab_size=32000, seq_length=4096, fp16_lm_cross_entropy=True, + logit_dtype=torch.float32, parallel_output=False, share_embeddings_and_output_weights=True, position_embedding_type="rope", @@ -699,6 +703,7 @@ def test_config_params_passed_to_mcore(self, mock_model, *_): assert kw["vocab_size"] == 32000 assert kw["max_sequence_length"] == 4096 assert kw["fp16_lm_cross_entropy"] is True + assert kw["logit_dtype"] == torch.float32 assert kw["parallel_output"] is False assert kw["share_embeddings_and_output_weights"] is True assert kw["position_embedding_type"] == "rope" diff --git a/tests/unit_tests/training/models/test_hybrid_builder.py b/tests/unit_tests/training/models/test_hybrid_builder.py index 9984e224ce3..fa304f7ca34 100644 --- a/tests/unit_tests/training/models/test_hybrid_builder.py +++ b/tests/unit_tests/training/models/test_hybrid_builder.py @@ -3,6 +3,7 @@ from unittest.mock import Mock, call, patch import pytest +import torch from megatron.core.transformer import ModuleSpec from megatron.core.transformer.transformer_config import TransformerConfig @@ -39,6 +40,7 @@ def test_builder_classvar(self): def test_default_values(self): config = HybridModelConfig(transformer=_make_transformer()) assert config.fp16_lm_cross_entropy is False + assert config.logit_dtype is None assert config.parallel_output is True assert config.share_embeddings_and_output_weights is False assert config.hybrid_layer_pattern is None @@ -55,6 +57,7 @@ def test_custom_initialization(self): config = HybridModelConfig( transformer=_make_transformer(), fp16_lm_cross_entropy=True, + logit_dtype=torch.float32, parallel_output=False, hybrid_attention_ratio=0.25, hybrid_mlp_ratio=0.1, @@ -63,6 +66,7 @@ def test_custom_initialization(self): vocab_size=50000, ) assert config.fp16_lm_cross_entropy is True + assert config.logit_dtype == torch.float32 assert config.parallel_output is False assert config.hybrid_attention_ratio == 0.25 assert config.hybrid_mlp_ratio == 0.1 @@ -329,6 +333,7 @@ def test_config_params_passed_to_mcore(self, mock_model, *_): seq_length=4096, hybrid_layer_pattern="M-A-", fp16_lm_cross_entropy=True, + logit_dtype=torch.float32, parallel_output=False, share_embeddings_and_output_weights=True, position_embedding_type="rope", @@ -345,6 +350,7 @@ def test_config_params_passed_to_mcore(self, mock_model, *_): assert kw["max_sequence_length"] == 4096 assert kw["hybrid_layer_pattern"] == "M-A-" assert kw["fp16_lm_cross_entropy"] is True + assert kw["logit_dtype"] == torch.float32 assert kw["parallel_output"] is False assert kw["share_embeddings_and_output_weights"] is True assert kw["position_embedding_type"] == "rope" diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index 93650cf13b0..9f6a20ae453 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -2,28 +2,37 @@ import gc +from unittest.mock import Mock, patch import pytest import torch from megatron.core import parallel_state from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedTensor +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.inference.contexts import StaticInferenceContext from megatron.core.inference.utils import InferenceMode from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_with_transformer_engine_submodules, ) +from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec from megatron.core.tensor_parallel.random import ( HAVE_TE, initialize_rng_tracker, model_parallel_cuda_manual_seed, ) -from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord -from megatron.core.transformer.enums import InferenceCudaGraphScope +from megatron.core.transformer.cuda_graphs import ( + CudaGraphManager, + _CudagraphGlobalRecord, + create_cudagraphs, +) +from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( + MoETransformerLayer, TransformerLayer, + TransformerLayerSubmodules, get_transformer_layer_offset, ) from megatron.core.utils import is_te_min_version @@ -378,6 +387,52 @@ def _no_layers_have_manager(block) -> bool: return all(not hasattr(layer, 'cudagraph_manager') for layer in block.layers) +def _make_moe_transformer_layer(*, partial_cudagraph: bool): + config = TransformerConfig( + num_layers=1, + hidden_size=32, + num_attention_heads=4, + ffn_hidden_size=64, + moe_ffn_hidden_size=64, + num_moe_experts=4, + moe_router_topk=2, + moe_router_load_balancing_type="none", + moe_token_dispatcher_type="allgather", + hidden_dropout=0.0, + attention_dropout=0.0, + bias_dropout_fusion=False, + add_bias_linear=False, + use_cpu_initialization=True, + cuda_graph_impl="local" if partial_cudagraph else "none", + cuda_graph_modules=[CudaGraphModule.moe_router] if partial_cudagraph else [], + ) + submodules = TransformerLayerSubmodules( + mlp=get_moe_module_spec(use_te=False, num_experts=4, moe_grouped_gemm=False), + mlp_bda=get_bias_dropout_add, + ) + return MoETransformerLayer(config, submodules) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_moe_router_synchronizes_host_outputs_and_reuses_event(): + layer = object.__new__(MoETransformerLayer) + host_output = Mock() + host_output.device.type = "cpu" + cuda_output = torch.empty(1, device="cuda") + event = Mock() + + with patch.object(torch.cuda, "Event", return_value=event) as event_factory: + layer._synchronize_router_host_outputs((host_output,)) + layer._synchronize_router_host_outputs((cuda_output,)) + assert event.record.call_count == 1 + assert event.synchronize.call_count == 1 + layer._synchronize_router_host_outputs((host_output,)) + + event_factory.assert_called_once_with() + assert event.record.call_count == 2 + assert event.synchronize.call_count == 2 + + @pytest.mark.skipif( not (HAVE_TE and is_te_min_version("1.5.0")), reason="CUDA graph tests require TransformerEngine >= 1.5", @@ -393,6 +448,41 @@ def teardown_method(self, method): _reset_cudagraph_state() gc.collect() + def test_moe_router_partial_cudagraph_forward_matches_eager(self): + eager_layer = _make_moe_transformer_layer(partial_cudagraph=False) + partial_cg_layer = _make_moe_transformer_layer(partial_cudagraph=True) + partial_cg_layer.load_state_dict(eager_layer.state_dict()) + eager_layer.cuda() + partial_cg_layer.cuda() + for param in partial_cg_layer.parameters(): + param.main_grad = torch.zeros_like(param) + + hidden_states = torch.randn(8, 2, 32, device="cuda", requires_grad=True) + eager_output, _ = eager_layer(hidden_states.clone(), attention_mask=None) + eager_output = eager_output.detach().clone() + + # The first forward/backward records the real router and postprocess graph boundaries. + recorded_output, _ = partial_cg_layer(hidden_states.clone(), attention_mask=None) + recorded_output.sum().backward() + create_cudagraphs() + + assert _CudagraphGlobalRecord.cudagraph_created + assert partial_cg_layer.use_partial_cudagraphs + for manager in ( + partial_cg_layer.cudagraph_manager_router, + partial_cg_layer.cudagraph_manager_postprocess, + ): + assert len(manager.cudagraph_runners) == 1 + assert manager.cudagraph_runners[0].fwd_graph is not None + + partial_cg_layer.zero_grad(set_to_none=True) + partial_cg_output, _ = partial_cg_layer(hidden_states.clone(), attention_mask=None) + partial_cg_output = partial_cg_output.detach().clone() + + # All-gather routing metadata stays on CUDA, so replay must not create a host-wait event. + assert not hasattr(partial_cg_layer, '_router_dtoh_event') + torch.testing.assert_close(partial_cg_output, eager_output, rtol=0, atol=0) + def test_empty_scope_transformer_layer_has_per_layer_manager(self): block = _make_cuda_graph_gpt_block( cuda_graph_impl='local', cuda_graph_modules=[], inference_cuda_graph_scope='layer' diff --git a/tests/unit_tests/transformer/test_utils.py b/tests/unit_tests/transformer/test_utils.py index 481a6bf5d61..9bbc97eb045 100644 --- a/tests/unit_tests/transformer/test_utils.py +++ b/tests/unit_tests/transformer/test_utils.py @@ -2,6 +2,7 @@ import inspect import os +from types import SimpleNamespace import pytest import torch @@ -13,11 +14,54 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import ( is_layer_window_attention, + set_model_config_attribute, set_model_to_sequence_parallel, ) from tests.unit_tests.test_utilities import Utils +class _TrackingConfig: + def __init__(self, value): + self._runtime_value = value + self.update_count = 0 + + @property + def runtime_value(self): + return self._runtime_value + + @runtime_value.setter + def runtime_value(self, value): + self._runtime_value = value + self.update_count += 1 + + +class _ConfigModule(torch.nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + + +def test_set_model_config_attribute_updates_distinct_configs_once(): + root_config = _TrackingConfig("original") + child_config = _TrackingConfig("original") + unsupported_config = SimpleNamespace() + + module = _ConfigModule(root_config) + module.first_child = _ConfigModule(child_config) + module.second_child = _ConfigModule(child_config) + module.unsupported_child = _ConfigModule(unsupported_config) + model = SimpleNamespace(config=root_config, module=SimpleNamespace(module=module)) + new_value = object() + + set_model_config_attribute(model, "runtime_value", new_value) + + assert root_config.runtime_value is new_value + assert child_config.runtime_value is new_value + assert root_config.update_count == 1 + assert child_config.update_count == 1 + assert not hasattr(unsupported_config, "runtime_value") + + class TestGPTModel: def setup_method(self, method): diff --git a/train_rl.py b/train_rl.py index acf54680f4a..17eec357b15 100644 --- a/train_rl.py +++ b/train_rl.py @@ -415,6 +415,10 @@ def _model_builder( extra_args_provider=add_inference_args, args_defaults={}, ) + assert not args.reset_attention_mask, ( + "--reset-attention-mask is not supported in RL training: " + "the forward pass masks via PackedSeqParams and never consumes a dense attention mask." + ) if is_hybrid_model(args): model_cfg = hybrid_config_from_args(args) else: