fix: correct LoRA initialization and forward pass under tensor parallelism

[chen2021673]

Summary

Fix LoRA loss divergence under tensor/data parallel training by making LoRA tensor-parallel initialization deterministic and aligning the forward collective order between base linear and LoRA linear paths.

This PR changes TP LoRA parallel linear modules to compute the base shard and LoRA shard locally first, add them before communication, and then run a single TP collective on the combined output. It also adds rank-aware Broadcast/Scatter ProcessGroup APIs and updates LoRA weight loading to support loading full saved LoRA tensors into TP-sharded model parameters.

It also adds rank-aware Broadcast/Scatter ProcessGroup APIs, updates LoRA adapter save/load to use portable full tensors under TP, handles packed QKV LoRA tensors, and keeps adapter IO compatible with DDP-wrapped models.

Motivation

In tensor-parallel LoRA training, the base linear output and the LoRA update together form one logical linear output. Previously, the base path and LoRA path could run separate TP collectives and add their outputs afterward:

base = TPCollective(base_i);
lora = TPCollective(lora_i);
out = base + lora;

This changes the floating-point communication/reduction order compared with computing the logical local contribution first:

out_i = base_i + lora_i;
out = TPCollective(out_i);

The separate-collective path can introduce numerical divergence across TP/DDP configurations.

Replicated or sharded LoRA parameters need rank-consistent initialization so every TP rank starts from the same logical LoRA weights.

Key Changes

LoRA parallel linear forward

• Update LoRAColumnParallelLinear to:

  • compute the base shard locally;
  • compute the LoRA shard locally;
  • add base_shard + scaled_lora_shard before TP gather;
  • run a single gather when gather_output_ is enabled.

• Update LoRARowParallelLinear to:

  • compute the base shard locally;
  • compute the LoRA shard locally;
  • add base_shard + scaled_lora_shard before TP reduce/reduce-scatter;
  • apply bias after the collective, matching the base RowParallelLinear behavior.

This makes the LoRA path follow the same logical communication boundary as the base parallel linear layer and avoids running separate collectives for base and LoRA outputs.

Deterministic TP LoRA initialization

• Make replicated ColumnParallel lora_A consistent across TP ranks.
• Initialize the logical RowParallel lora_A from TP rank 0 and distribute the correct local shards to each TP rank.
• Ensure TP ranks observe the same logical LoRA initialization instead of independently sampling incompatible parameter values.

ProcessGroup communication APIs

• Add rank-aware communication APIs:

  • ProcessGroup::Broadcast(tensors, root_rank_in_group, async_op)
  • ProcessGroup::Scatter(output_tensors, input_tensors, root_rank_in_group, async_op)

• Rename the old APIs to BroadCast_ / Scatter_ to avoid semantic ambiguity.

• Update existing autograd communication wrappers to call the legacy renamed APIs where appropriate.

LoRA weight save/load

• Extend LoadLoRAWeights to support loading full saved LoRA weights into TP-sharded model parameters.

• When a destination parameter is TP-sharded, slice the loaded full tensor according to the current tp_rank before copying it into the local parameter.

• Update SaveLoRAWeights to export portable full LoRA weights under TP:

  • TP-sharded LoRA tensors are gathered back into full/unsharded tensors before writing.
  • All TP ranks in the primary DP/PP group participate in the save-side gather.
  • Only the rank with dp == 0 && tp == 0 && pp == 0 writes the weight file.

Packed QKV LoRA tensor handling

• Handle attention c_attn.lora_B specially when it uses the packed QKV layout [Q | K | V].
• On save, restore full packed QKV tensors from per-rank [Qi | Ki | Vi] shards.
• On load, split full packed QKV tensors back into local [Qi | Ki | Vi] shards for the current TP rank.
• Add helper utilities for packed QKV slicing/restoring and unit tests covering GPT-style and GQA-style QKV layouts.

DDP adapter state dict compatibility

• Expose the wrapped module from DistributedDataParallel and save LoRA adapters from the inner model.
• This matches the common model.module.state_dict() behavior for DDP-wrapped models.
• Consume a legacy leading module. prefix when loading LoRA adapter state dicts.
• Add test coverage for DDP module. prefix compatibility.

Test

The loss values in the LoRA tests fluctuate slightly, which is expected because the tests now use saved fixed initialization values.

The performance variance is concentrated in the lora/bfloat16/distopt tests and is treated as normal fluctuation.

[Chamberlain0w0]

现在模型实现中，Attention 里面第一个 Linear 把 QKV 合并了（实际排列上是 [Q | K | V]），LoadFromLLMC 里面有对应的处理（如 TP=4，则切成 [Q0 Q1 Q2 Q3 | K0 K1 K2 K3 | V0 V1 V2 V3]，然后依次拼接 [Qi | Ki | Vi] load 到每个 tp rank 上）。所以这块的切分逻辑对于这种情况没法适用，得想想看怎么能特别处理一下

[chen2021673]

加一个特判，在 LoadLoRAWeights 里识别 attn.c_attn 这种输出维分片参数，把 full [Q|K|V] 切成 local [Qi|Ki|Vi]。

[chen2021673]

checkpoint 没单独处理这个类型，load 的时候必须对得上 TP 的配置。但这里因为要和单机对齐初始化参数，有这个功能需求。

[kilinchange]

麻烦修改一下 3b6088e 这个 commit 的 message，然后再跑一遍全量测试吧。

[chen2021673]

麻烦修改一下 3b6088e 这个 commit 的 message，然后再跑一遍全量测试吧。

已修改。
测试截图已更新到PR描述。