peft: fix multi-LoRA activation-recompute replay and grouped-expert adapter DDP flagging

[yushengsu-thu]

Two one-line-class fixes for silent multi-LoRA training failures on MoE models, found by GPU triage on the miles tinker stack (GPT-OSS 20B, expert-only LoRA, 4xH200).

Fix 1 — activation-recompute replay never happens for multi-LoRA (peft/recompute.py)

maybe_enable_recompute_inputs_grad ships the input-grad mechanism that makes activation-checkpoint replay work for frozen-base PEFT (a checkpointed region is only replayed grad-enabled when its input requires grad). Its name matching only recognized single-LoRA params (.adapter.); multi-LoRA slot params live in an adapters nn.ModuleList (.adapters.<slot>.) and were classified as trainable base weights, so the patch was skipped. Under --recompute-granularity full (and selective moe with expert-only targets) the checkpointed region gets grad-free inputs, CheckpointFunction.backward is never invoked, and every adapter gradient is identically zero — silent no-op training at a truthful grad_norm=0.0.

Fix: extend the matching to .adapters. via a shared helper. Single-LoRA classification is byte-identical.

Fix 2 — grouped-expert adapter grads never reduced across TP (peft/utils.py)

GroupedExpertLinearAdapter.__init__ set allreduce = not (EP > 1) and never set tensor_model_parallel. With EP=1, ETP=1, TP>1 the TP-duplicated adapter weights land in the dense DDP bucket (reduced over dp_cp only), so cross-TP gradient reduction never happens: rank>0 grads stay identically zero and the TP replicas of the adapter silently diverge from step 1.

Fix: allreduce=False unconditionally. Expert adapter weights are sharded over ETP (not TP), so their grads must be reduced over the expert data-parallel group, which folds in the TP peers whenever ETP < TP; when ETP == TP the expert-DP group equals dp_cp, so previously-working configs (including EP > 1) are unchanged. Comment states the invariant.

GPU evidence (4xH200, GPT-OSS 20B bf16 expert-only LoRA via miles tinker stack @ tinker-frontend d5bb5e062; TP=2+SP, EP=1, ETP=1, 2 adapters)

Check	before (bridge @ 3fc31b0)	after (this branch)
--recompute-granularity full: SDK grad_norm per step	0.0 on every step, both adapters (silent no-op)	45.07 / 29.13 (rank-8 adapter), 119.21 / 30.08 (rank-16)
--recompute-granularity full: trainer logprob max-abs delta after optim_step	no real learning signal	1.01 / 2.25 and 2.45 / 1.36
adapter forward grad-enabled during checkpoint replay	never	yes — e.g. 96 grad-enabled wrapper forwards per step (replay pass) vs 96 no-grad (checkpoint pass)
TP=2 adapter main_grad across ranks (sha256, post-reduce pre-step)	rank1 identically zero on all 96 params, every step; replicas fully divergent (0/96 hashes match by step 2)	bit-identical: 96/96 param hashes match on every optim step, both slots, across all runs (17 optim calls total)
poison-window 5-phase GPU acceptance (tinker_sdk_poison_window.py)	strict probe-determinism precondition already unmet on this MoE config: back-to-back forwards differ by 0.16-0.21 max-abs-dlogprob with NO optim step (base bridge, probe_stability diagnostic)	same pre-existing forward nondeterminism (0.09-0.14 — not a regression); all 5 phases PASS with every mechanism assertion kept and probe/grad-norm tolerances sized to the measured noise: typed rejects in 0.1s, clocks held, discard leaves no residue (recovery grad_norm 157.53 vs ref 157.96 and 151.12 vs 151.19), isolation clean, late-chunk discard held

Logs and dumps persisted on cluster h200-sci-k8s under /personal/bridge-fix-0812/ (prior failing-run evidence: /personal/triage-positive{,2,3}, passing no-recompute baseline: /personal/triage-positive4).

Tests

• tests/unit_tests/peft/test_recompute.py: multi-LoRA (.adapters.<i>.) params are recognized (regression: patch used to be skipped); trainable-base early-out still holds; existing single-LoRA test unchanged.
• tests/unit_tests/peft/test_utils.py: GroupedExpertLinearAdapter weights always carry allreduce=False + correct partition_dim/partition_stride (parametrized over input_is_parallel), with EP=1 in the config — the exact regression scenario.
• Full tests/unit_tests/peft/ on the H200 box: 362 passed, 9 pre-existing failures identical at base commit 3fc31b0 (env-related, files untouched by this PR).

Note for miles

miles' launch-time guard refusing full recompute (radixark/miles#2273, commit c0def7864) was added assuming an unfixed bridge; it can be relaxed to a bridge-version check once this merges. The guard is not yet on tinker-frontend, so validation ran unmodified miles @ d5bb5e062.

[yushengsu-thu]

Reworked the utils.py fix per request (afd22d6): the unconditional allreduce=False is reverted and replaced with a port of the upstream (NVIDIA-NeMo/Megatron-Bridge main) form — allreduce = not use_expert_process_groups where use_expert_process_groups = EP > 1 or ETP != TP, plus set_tensor_model_parallel_attributes(weight, True, tp_axis, 1) (which also sets tensor_model_parallel). Outcomes are identical in every configuration (when the condition is false, the expert-DP group equals dp_cp), so the existing GPU evidence (EP=1, ETP=1, TP=2 → allreduce=False → bit-identical cross-rank grads) remains valid. The regression test now covers both sides of the condition: ETP≠TP → expert bucket; ETP==TP → dense (equivalent groups).