Repository navigation
Context parallelism for full fine-tuning (H100 128K) - #8
Merged
Merged
Conversation
ApGa
commented
Oct 14, 2026
ApGa
commented
Collaborator
- Ports upstream #3320's CP loss layer and adds CP gathers for Tinker forward logprobs and per-datum outputs.
- Fixes a silent gradient bug. At CP>1, megatron-bridge's Qwen3.5/3.6 VL providers turn per-token loss back on after miles disables it, which scales every gradient by microbatches × DP·CP. The flag is now reset after provide(), and the DDP scaling factors are asserted.
- Guards:
- CP>1 is allowed for full fine-tuning only, with thd packing and a CP-aware datum cap.
- R3 is rejected together with router fusion.
- Startup fails if the MRoPE patch is missing.
- GDN layers default to headwise CP.
- Tooling and docs: launcher knobs (--cp, --linear-cp-mode), a Qwen3.5-4B recipe, the parity client deploy/clients/tinker_cp_parity.py with CPU and GPU probes, and documentation.
- Validation, Qwen3.5-4B full FT on ABCI:
- CP2 (headwise) vs CP1 gives a step-1 grad-norm ratio of 0.9993.
- Forward logprob deltas are at the TP2 noise level, with no chunk-boundary errors.
- Weight-update equality, the SDK smoke (including save and resume) and on-policy KL pass.
- CP4 and CP8 grad norms are within 0.2%.
- Not yet validated: 27B and 35B at 128K on H100. The planned layouts are TP4×CP2×DP4 and TP1×CP8×EP8×DP4.
Ported from radixark/miles 3d8b528 (#3320) onto the fork's tinker_losses.py and its per-datum recorder. At CP>1 each rank slices the full-length weights, advantages and masks to its zigzag shard (the padding gate still reads the full mask), the loss stays a sum over the local shard, and a rank with no response token returns logits[..., :0].sum() so it still runs the shared backward. A standalone helper in cp_utils.py rebuilds the per-datum logprobs and losses with one CP all-reduce per microbatch, so every rank reports full-length outputs. Gloo tests cover CP 1/2/4 with gateway-shaped and tiny datums, pad rows, pad-only microbatches, the fused head and a recomputed loss.
forward_only_logprobs returned each CP rank's zigzag shard of every datum and the gateway kept whichever copy arrived first. It now rebuilds full-length logprobs with the same batched gather as the per-datum recorder (one CP all-reduce per unit), after the R3 replay stage and the fused head have run.
The Qwen3-VL and Qwen3.5-VL bridge providers turn calculate_per_token_loss on in finalize() at CP>1. Megatron's DDP and schedule read that config while miles' loss scaling reads args, so every gradient came out num_microbatches * DP*CP too large. The bridge provider now resets the flag on every config of the built model, and on the routers' cached copy, after the bridge's own assert. After get_model, a check verifies the config flag and that every DDP grad buffer scales by 1/(DP*CP). At CP>1 a bridge Qwen3VLModel now refuses to build without miles' packed MRoPE patch. --cp-comm-type and a new --linear-cp-mode (GatedDeltaNet layers) are passed on to the bridge provider. With R3 on, the router must not be fused.
The Qwen3.6 gateway launcher gains --cp and --linear-cp-mode. At CP>1, --max-tokens-per-gpu is a per-rank budget and defaults to the zigzag shard of one context_len + 1 row (66048 at TP4/CP2, 33280 at TP4/CP4, 16512 at TP1/CP8 for 128k rows). A new 4b variant serves Qwen3.5-4B, a small bridge VL model with GDN layers, with the standard head (its embeddings are tied) and TP1 engines. Serve validation now checks --tinker-max-tokens-per-datum against the per-rank budget (#3320's cap: an error at CP>1, a warning at CP=1 only above the budget) and rejects R3 with a fused MoE router.
deploy/clients/tinker_cp_parity.py run sends a fixed, seeded set of datums to a gateway: 2- to 17-token rows (some CP rank holds one or no target), rows up to 8K tokens cut from a repeated English passage, and a lone datum that leaves some DP ranks with pad-only microbatches. It records Tinker forward logprobs, the per-datum losses and logprobs of cross_entropy, importance_sampling and cispo (one per optimizer step), the three grad_norms and the logprobs after the steps. compare reports mean, p99 and trimmed-KL logprob deltas overall, per zigzag chunk and near chunk boundaries, plus loss deltas and grad_norm ratios, and exits 1 on a threshold failure. --noise takes the dump of a correct alternative layout (TP2 at CP=1) and widens the limits to its deviation. deploy/abci/clients/cp_parity.sh runs it as a CLIENT_SCRIPT, followed by the on-policy probe and the SDK smoke.
docs/fork/full-finetuning.md gains a CP section: what the gateway does under CP (losses, gathers, the bridge per-token-loss reset, MRoPE, GDN modes, R3), the knobs and per-rank token budgets, the limits, and the Qwen3.5-4B validation on H200. There the step-1 grad_norm matches CP=1 within 0.23% at CP2 and 0.19% at CP4, and the forward deltas match a TP2 run. FORK.md lists the CP files.
The Tinker gateway accepted only --context-parallel-size 1. Full fine-tuning now runs at CP>1, now that the losses, both gathers and the bridge fixes handle the zigzag shards. CP>1 requires --qkv-format thd and rejects --allgather-cp (the losses, gathers, bridge pre-shard detection and R3 fill all assume the zigzag layout), and the per-datum token cap is checked per rank. The LoRA gateway stays at CP=1.
apply_context_parallel_modes now runs only at --context-parallel-size > 1. Megatron's default --cp-comm-type p2p no longer overrides a provider's own choice (megatron-bridge's Mistral provider picks a2a for sliding windows), and --linear-cp-mode is unset by default, so native and bridge configs keep TransformerConfig's chunkwise default and a provider without the field still builds. The post-get_model check now asserts what DDP's constructor does not: each chunk reduces over a DP*CP group the size of the intra_dp_cp factor loss.py multiplies by. full/executor.forward_backward re-checks the per-token-loss flag that the schedule reads on every call. The docstring corrects the impact of the bridge flag. Only the gateway's deferred finalize kept the num_microbatches * DP*CP overscale. Elsewhere finalize_model_grads cancels it for the main loss, and the reset changes the MoE aux/z-loss and MTP loss scales instead.
The startup check proved the packed-MRoPE patch was installed, not that a forward used it. On a CP-sharded THD row the patched Qwen3VLModel.forward now raises when no rank-local positions were built (the bridge's dense get_rope_index would number the zigzag chunks as one sequence), and when the bridge never asked get_rope_index for them.
Every CP and TP rank of a DP group holds the same gathered vectors after the CP all-reduce, and each put a copy into the object store. Only CP/TP rank 0 returns them now; the runtime already skips the empty results.
The chunkwise GDN backward is unvalidated: on natural text the step-1 grad_norm moved 0.2% at CP2 against 0.02% for TP2, and 24% on random motifs. At --cp > 1 the launcher now passes --linear-cp-mode headwise unless chunkwise is requested, and checks that TP*CP divides the recipe's linear key/value heads (16/32 for the 4B and 35B, 16/48 for the 27B). The per-rank budget pads with the --data-pad-size-multiplier that --extra-args sets, not a fixed 128.
Chunkwise GDN CP hands its state across a contiguous split of the whole packed row, not at the zigzag boundaries the probe checked, and short datums were diluted among the 8K-token rows. run now sends a lone 4000- and 8000-token datum (forward, and forward_backward in the first window with the loss on the second half) and records TP and the pad multiplier. compare gates windows at the split rows of those datums, roundup(2*ceil(L/2CP), TP*pad) apart, gates each short row by its max |dlogprob| before the first optimizer step, and keeps short rows out of the KL trim.
The replay rows fill_replay_data registers must follow get_batch's zigzag tokens on every CP and TP rank. A gloo-free test now checks CP 1-8, TP/SP 1-4 and three pad multipliers on gateway-shaped datums, plus a negative control (a contiguous routing slice misaligns at CP>1). R3 at CP>1 is still unchecked on GPU, so serve validation warns unless --tinker-r3-check is on. Also restores the original formatting of an unrelated fused-head assert.
The gateway keeps rank 0's metrics, but ranks differ under R3 load
imbalance and uneven CP shards. Each window now also reports the maximum
over all trainer ranks (one gloo all-reduce) as
actor_peak_mem_{alloc,reserved}_gib_all_ranks.
tools/fork/gdn_cp_grad_probe.py runs one Megatron GatedDeltaNet layer (Qwen3.5-4B sizes) under torchrun on a packed THD row shaped like get_batch's, without CP and with each linear_cp_mode at CP 2/4/8, from the same weights and inputs. The loss sits after the first contiguous CP split of the datum that straddles it, and half the heads decay slowly, so the pre-split input gradients and part of every parameter gradient cross ranks. It prints the cosine and norm ratio of each against the no-CP reference.
--layers attention runs a Qwen3.5-4B full-attention layer (TE flash attention, qk layernorm, output gate, no rotary) through the same comparison, once per --attention-modes cp_comm_type. The grad_norm shift under CP2 shows with both GDN modes, so the attention ring is the other candidate.
The zigzag windows sit near the start of the short rows, which differ more under any layout change: on natural text a correct TP2 run already shows 3.1x the interior delta there, and headwise CP2, quieter in the interior, reached 3.7x against the fixed 3x. With --noise, a window may now reach 1.5x the window/interior ratio the noise run shows at the same rows. This applies to the GatedDeltaNet split windows too, instead of an absolute widening.
On the Qwen3.5-4B natural-text dumps a correct TP2 run and both GDN modes at CP2 reach 3.7x the interior delta in the zigzag windows, so the default ratio is now 4 (1-token boundary errors are still caught from 0.1 nats). The lone 4000/8000-token rows average ~4000 tokens of repeated text, where a few flips move their loss by percents under any layout (TP2: 5%), so their loss is reported, not gated, and their aggregate deltas are gated only against a noise run.
At initialization the qk-normalized logits have unit variance, so each query spreads over thousands of keys and bf16 rounding of the nearly cancelling partial results dominates the comparison. --attention-sharpness scales the q/k layernorm gains to a chosen logit standard deviation.
The attention layer under ring attention disagreed with no CP far more than bf16 noise (output cosine 0.90 with sharp logits). The probe now counts the rows below cosine 0.99 per segment and zigzag chunk, builds the attention without linear biases as Qwen3.5 does, and can drop the output gate (--no-attention-gate) to isolate it.
TransformerConfig defaults to attention and hidden dropout of 0.1, which made every attention output row differ between layouts (the gateway trains with both at 0). The GatedDeltaNet layer has no dropout and was unaffected.
docs/fork/full-finetuning.md now states the bridge per-token-loss impact correctly (only the gateway's deferred finalize kept the gradient overscale; elsewhere it changed the MoE aux/z-loss and MTP scales), the headwise launcher default and its head constraint, the new MRoPE, R3 and loss-scaling checks, the all-rank peak memory, and how the parity probe covers GatedDeltaNet splits and short rows. The validation table covers CP2/4/8 headwise and CP2/4 chunkwise against CP1 with TP2 as the noise run (all pass), and the layer probes (GDN in both modes and ring attention match no CP at bf16 level). It records the open item: step-1 grad_norm is 0.16-0.30% above CP1 at every CP size in both modes, which neither layer type explains. FORK.md and the ABCI README follow.
CP2 headwise (job 2635004) matches a CP1 reference packed at the same 8448 tokens per microbatch (job 2635003): compare passes without a noise run, the step-1 grad_norm ratio is 0.9993, KL(sampler, trainer) is 2.8e-4, the SDK smoke has 0 failed checks and every engine passes the weight check. The 0.16-0.30% step-1 grad_norm offset left open in round two came from the reference's packing, not from CP: CP1 alone moves 0.31% between 8193 and 8448 tokens per microbatch. The parity instructions now give the CP1 reference the candidate's row budget.
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
1 participant
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.