[bug] LoRA adapter init depends on tensor parallel size: linear_in is initialized per shard with a fan-dependent method #5840

Problem

LoRA training results depend on tensor parallel size, because the adapter's A matrix is initialized per TP shard with a fan-dependent init method.

For LoRA on a row-parallel base layer (linear_proj, linear_fc2), ParallelLinearAdapter builds linear_in as a RowParallelLinear, so the A matrix is sharded along its input dimension. Megatron-core initializes each shard locally (_initialize_affine_weight_gpu applies the init function to the shard tensor). The default lora_A_init_method="xavier" -> nn.init.xavier_normal_ computes its std from the tensor it is applied to — the shard, not the full matrix. So the A init std is sqrt(2 / (in_features/TP + dim)) instead of sqrt(2 / (in_features + dim)): it grows roughly with sqrt(TP), and the adapter trained at TP=4 starts from a different distribution than the same config at TP=1.

Measured (Qwen3-8B, rank 32, seed fixed):

Training the same data with the same recipe, the converged loss moves monotonically with TP width: the mean NLL over the last epoch of a 222-step LoRA run shifted by 0.013 from TP1 to TP4 and 0.021 from TP1 to TP8, far outside run-to-run noise (~1e-4). DP size and GPU count had no effect; only TP moved it. Confirmed as the cause by intervention: computing the same xavier std from the full tensor fan (sqrt(2/(in_features + dim)), passed as a fixed-std normal init) makes results TP-invariant — TP4 and TP8 then land within run-to-run noise of TP1.

Root cause: Megatron-core's per-shard GPU init is only distribution-safe for fan-independent init functions (base model weights use init_method_normal(sigma) with a precomputed sigma for exactly this reason). The PEFT path passes fan-dependent functions (xavier, also kaiming) into that machinery. "normal" and "zero" are unaffected.

Possible fix: in ParallelLinearAdapter.__init__, resolve fan-dependent init methods to a concrete std from the full (unsharded) dims — e.g. map "xavier" to a fixed-std normal with sqrt(2/(in_features + dim)) before passing it to the parallel linear. Same for linear_out when it is sharded and a fan-dependent method is configured.

Minimal repro

Apply LoRA(target_modules=["linear_proj", "linear_fc2"], dim=32) (defaults otherwise) to any model at TP=1 and at TP=4, and compare adapter.linear_in.weight.std() — expect ~2x. Training the same data at both TP sizes shows diverging loss curves.

Expected behavior

Tensor parallel size does not change the initialization distribution (or any other training semantics) of the adapter.

Affected area

area:peft

Regression?

No — long-standing behavior of the per-shard init + fan-dependent default combination, not tied to a recent refactor as far as we can tell.

Environment

Megatron-Bridge at abeec217, H200, bf16, GPU init (use_cpu_initialization=False). Also reproduces through consumers of this path (e.g. NeMo RL's Megatron policy uses this same LoRA class).

Logs

Init-time std measurements and the TP1/TP4/TP8 loss trajectories are available on request (measured with a fixed seed and a frozen 222-step data stream; every number above reproduces run-to-run).