[Dtensor] SFT with LoRA is slower than without LoRA

[RayenTian]

Describe the bug

When enabling LoRA on the DTensor backend, end-to-end training throughput degrades compared to the same setup without LoRA. This appears even when LoRA rank is small and other settings are unchanged.

Reproduce

Disable Lora

NRL_FORCE_REBUILD_VENVS=true uv run examples/run_sft.py \
logger.wandb_enabled=True \
logger.wandb.project=lora \
logger.wandb.name=nano_v3_lora \
policy.model_name=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
checkpointing.enabled=False \
policy.dtensor_cfg.enabled=true \
policy.dtensor_cfg._v2=true \
policy.dtensor_cfg.lora_cfg.enabled=False \
policy.dtensor_cfg.lora_cfg.use_triton=False \
policy.dtensor_cfg.lora_cfg.dim=8 \
policy.max_total_sequence_length=2048 \
policy.train_global_batch_size=16 \
policy.train_micro_batch_size=1 \
policy.optimizer.name="torch.optim.Adam" \
~policy.tokenizer.chat_template  \
sft.max_num_steps=10 \
cluster.num_nodes=2 \
cluster.gpus_per_node=8

Enable Lora

NRL_FORCE_REBUILD_VENVS=true uv run examples/run_sft.py \
logger.wandb_enabled=True \
logger.wandb.project=lora \
logger.wandb.name=nano_v3_lora \
policy.model_name=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
checkpointing.enabled=False \
policy.dtensor_cfg.enabled=true \
policy.dtensor_cfg._v2=true \
policy.dtensor_cfg.lora_cfg.enabled=True \
policy.dtensor_cfg.lora_cfg.use_triton=False \
policy.dtensor_cfg.lora_cfg.dim=8 \
policy.max_total_sequence_length=2048 \
policy.train_global_batch_size=16 \
policy.train_micro_batch_size=1 \
policy.optimizer.name="torch.optim.Adam" \
~policy.tokenizer.chat_template  \
sft.max_num_steps=10 \
cluster.num_nodes=2 \
cluster.gpus_per_node=8

Observed Behavior

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🔗 Related PRs

#1556 - feat: LoRA SFT support for DTensorV2 path [merged]

👤 Suggested Assignees

• samodi-nv

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord or schedule a call!

[terrykong]

@RayenTian are you looking into this? i imagine this slowdown will translate to GPRO