Wire witness support inside Megatron - #21
Merged
Merged
Conversation
Add _pre_decoder_hooks list and register_pre_decoder_hook() method. Hooks are called between _preprocess and decoder, allowing external code to transform decoder_input without Megatron knowing specifics.
Already initialized in __init__, no need for getattr fallback.
fzyzcjy
marked this pull request as draft
April 2, 2026 02:28
Remove _pre_decoder_hooks list and register_pre_decoder_hook() method. Add witness_ids parameter to forward() and build_schedule_plan(). Witness logic is inline: hasattr(self, 'head_witness') check + add to decoder_input or decoder.input_tensor depending on PP stage.
- build_schedule_plan: accept witness_ids, pass to schedule plan - TransformerModelChunkSchedulePlan: store witness_ids in chunk_state - PreProcessNode.forward_impl: apply witness after _preprocess (same logic as GPTModel.forward)
- GPTModel.forward: add tail_witness after decoder, before _postprocess - build_schedule_plan: revert witness_ids param (not supported) - model_chunk_schedule_plan: revert witness_ids in chunk_state - fine_grained_callables: revert witness logic in Pre/PostProcessNode
_DataWitness.forward returns [b, s, 1] but Megatron's decoder_input is in [s, b, h] format after the embedding layer. Without transposing, broadcasting [s, b, h] + [b, s, 1] creates a [s, s, h] tensor, causing OOM (648 GiB for s=18432).
When sequence parallel is active, decoder_input and hidden_states are scattered along the sequence dimension ([s/tp, b, h]). The witness output must also be scattered to match, otherwise shapes mismatch (e.g. [18432, 1, 1] vs [9216, 1, h] with TP=2).
…ptimizer The distributed optimizer replaces optimizer param_groups with shard main params (fp32 copies). get_main_grads_for_grad_norm checks _is_witness_param on these main params, but the flag was only set on the original model params. Copy the flag when building main param groups for both float16→fp32 and fp32 paths.
…ixes) Brings in: true-on-policy init (mappings.py _tree_all_reduce_sum + _DeterministicReduceFromModelParallelRegion, miles_megatron_plugins/true_on_policy, gpt_layer_specs use_true_on_policy_backend) + plugin packaging (#49, #50) + fixes (param dtype before DDP wrap #24, attention_output_gate TP slice #22, fp32 shard detach).
Debug/test-only helper that replaces order-sensitive SUM collectives with an all-gather plus a fixed local fold. For power-of-two world sizes it reuses the pairwise tree of _tree_reduce_sum_from_gathered; non-power-of-two falls back to an ascending-rank sequential fold. The all-gather is pure data movement, so the local sum order is the only arithmetic order, making the result bitwise deterministic and independent of NCCL.
When deterministic collectives are enabled, start_grad_sync folds each bucket's full grad_data with deterministic_sum_inplace over the same communication group the normal path selects, then returns early (skipping _coalescing_manager). The full reduced sum lands in grad_data on every rank, a superset of reduce-scatter, so the distributed optimizer's shard views read correct values. Synchronous semantics only: asserts overlap_grad_reduce is off and a single DistOpt instance; grad_reduce_handle stays None as in the existing non-overlap path.
Gate every order-sensitive SUM all_reduce that feeds the deterministic FT comparison behind is_deterministic_collectives_enabled(), preserving the original call on the else branch: - finalize_model_grads: conditional-embedding (pp), embedding (embd), non-tensor-parallel SUM grads (tp; the AVG entry is left untouched), and num_tokens (dp_cp). - moe_utils.get_updated_expert_bias: the router tokens-per-expert reduce (local import to avoid a circular import with megatron.core.distributed). - clip_grads.get_grad_norm_fp32: the two squared-norm SUM reduces; the inf-norm MAX reduce is left untouched.
Cover: deterministic_sum_inplace equals the true sum, bitwise-identical results across two process-group objects over the same ranks, chunked folding matching a single-shot fold, power-of-two tree vs non-power-of-two ascending fold selection, world-size-one no-op, and the contiguous-input assertion. The all_gather is faked in-process, so the test runs without GPUs.
Reuse _is_power_of_two from tensor_parallel.mappings; make fold_gathered_sum public (consumed by miles indep_dp) and let it reuse the first gathered buffer instead of clone+reallocate in the ascending fold; extract deterministic_sum_inplace_with_gather (injectable all_gather) so non-c10d process groups share the exact chunk/fold code; drop the redundant .contiguous() on contiguous 1-D slices; comment the SUM-only gate next to the AVG entry; cover the injectable-gather variant in the unit test.
…iguous input, group-ranks log) M1: the shared _patch_all_gather fake copied the full per-rank tensor into each sub-tensor chunk, which crashes for chunked folds. Slice by a running per-call offset like _fake_gather, so chunked and single-shot tests both work. m1: deterministic_sum_inplace_with_gather previously asserted contiguity, but finalize_model_grads' embedding-grad path may pass non-contiguous tensors that plain all_reduce accepted. Fold a contiguous copy and write back via copy_; never fall back to NCCL all_reduce. Contiguous fast path stays allocation-free. The unit test now asserts the correct summed values instead of an assertion. I2: add log_group_ranks_once helper and call it from the gated branch in start_grad_sync so each run logs the actual dp x cp member ranks once, enabling in-pod cp-inner bracketing verification.
…l backend) Reverts c30d54d, 336dd33, 02b995b, 85c2ce9, bb8aa73, 3998005. The deterministic SUM fold now lives entirely in miles as a custom torch.distributed backend (miles.utils.det_process_group, registered as det_nccl), so Megatron needs no gated call sites, no deterministic_collectives module, and carries zero extra divergence beyond upstream miles-main.
When MILES_EXPERIMENTAL_FT_TRAINER=1, assert CUDA_DEVICE_MAX_CONNECTIONS != 1 instead of == 1: the FT (indep_dp) trainer deadlocks a freshly respawned cell's concurrent intra-cell comms on rejoin under a single hardware queue.
…nt trainer" This reverts commit 07d03a1.
fzyzcjy
marked this pull request as ready for review
June 17, 2026 11:04
fzyzcjy
changed the title
fzyzcjy
changed the title
fzyzcjy
changed the title
yueming-yuan
approved these changes
Jul 9, 2026
Zhichenzzz
approved these changes
Jul 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.