[peft] fix: Add replica_id handling for dense LoRA adapters with TP > 1

[yaoyu-33]

Summary

When using LoRA adapters on dense layers (non-MoE) with TP > 1, only one TP shard was being saved during checkpointing. This caused significant training-inference loss discrepancy because other TP ranks loaded zero/uninitialized adapter weights.

Root Cause

In ParallelLinearAdapter.sharded_state_dict(), there was special replica_id handling for expert adapters (added in PR #1564, related to verl-project/verl#4303), but dense adapters never received the equivalent fix.

For dense adapters with TP > 1:

• Each TP rank generates a ShardedTensor with the same replica_id
• Shards should be distinguished by global_offset, but during the PEFT-filtered checkpoint save, TP shards are incorrectly deduplicated
• Result: Only TP rank 0's shard is saved

Solution

Add replica_id handling for dense adapters similar to expert adapters. When TP > 1, the replica_id is adjusted to include the TP rank, ensuring each TP shard is correctly identified and saved during PEFT-filtered checkpoint saves.

Changes

• Add replica_id adjustment for dense adapters with TP > 1 in ParallelLinearAdapter.sharded_state_dict()
• Add unit tests for the fix covering:
  • Dense adapters with TP > 1 (replica_id correctly updated)
  • Dense adapters with TP = 1 (no change to replica_id)
  • Expert adapters (still use EP-based replica_id, not affected)

Testing

• Unit tests added
• Functional tests with multi-GPU TP > 1 (requires CI)

Summary by CodeRabbit

• Bug Fixes

  • Fixed replica ID assignment for adapter checkpoints when using tensor model parallel configurations, ensuring correct deduplication during checkpoint saving.

• Tests

  • Added comprehensive test coverage for dense and expert adapter scenarios across different parallel configurations.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test acaef1a

[coderabbitai]

📝 Walkthrough

Walkthrough

This pull request modifies PEFT state dict handling for dense adapters with tensor model parallel (TP) size greater than 1, adjusting replica_id values for correct checkpoint deduplication. Comprehensive test coverage is added for parallel linear adapter state dict scenarios across different parallelism configurations.

Changes

Cohort / File(s)	Summary
PEFT State Dict Logic src/megatron/bridge/peft/utils.py	Removed in-function import of parallel_state and added new else branch for dense adapters with TP > 1 to adjust replica_id to (original_dim0, tp_rank, original_dim2) across linear_in_sd and linear_out_sd for correct deduplication when saving PEFT-filtered checkpoints.
Test Coverage tests/unit_tests/peft/test_utils.py	Added three new test cases for ParallelLinearAdapter sharded state dict handling: one for dense adapters with TP > 1 verifying replica_id includes TP rank, one for TP = 1 verifying no changes, and one for expert adapters verifying EP-based calculation is unaffected by dense adapter fix.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 62.50% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Test Results For Major Changes	✅ Passed	PR qualifies as minor changes: small focused code change (+14 lines) with comprehensive unit tests (+157 lines) covering key scenarios, though functional tests pending in CI.
Title check	✅ Passed	The title clearly and concisely describes the main change: fixing replica_id handling for dense LoRA adapters when TP > 1, which is the primary purpose of the changeset.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch fix/dense-lora-adapter-tp-replica-id

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[yaoyu-33]

/ok to test 8094b2b

[priyatham-resolve]

LGTM. Thanks for the fix

[priyatham-resolve]

nit: This seems to duplicate all the helper methods from TestLoRAFinetune above. Could share them via a base class or mixin to reduce the ~300 lines of boilerplate? Not a blocker.