[bugfix]: synchronize post-FSDP LoRA replicas

[Suckl]

Purpose

Fix replicated LoRA training state in the new YAML/modular training stack when LoRA adapters are injected after FSDP/HSDP wrapping.

These late-added parameters are not managed by FSDP gradient-reduction hooks. The LoRA forward path converts their replicated DTensors to local tensors, so rank-local gradients can be labeled as Replicate() without actually being synchronized. This can cause different optimizer updates, gradient clipping coefficients, and LoRA weights across ranks.

This PR is limited to the new modular training stack. The legacy training stack and general checkpoint strictness are out of scope.

Changes

• Broadcast late-added LoRA parameter values when creating replicated DTensors, ensuring the Replicate() placement matches the actual data on every rank.
• Register post-accumulate gradient hooks for replicated LoRA parameters.
• Average gradients across every replicated FSDP/HSDP mesh dimension before gradient clipping and the optimizer step.
• Preserve correct behavior across multiple backward passes with gradient accumulation.
• Add two-rank distributed regression coverage for:
  • initial parameter synchronization;
  • rank-local gradient averaging;
  • gradient accumulation;
  • identical optimizer updates across ranks;
  • CPU/Gloo and CUDA/NCCL execution.

Test Plan

# New LoRA distributed tests
python -m pytest fastvideo/tests/train/utils/test_lora.py -v

# Complete modular training utility suite
python -m pytest fastvideo/tests/train/utils -q

# Full repository checks on Linux
pre-commit run --all-files

# Two-GPU CUDA/NCCL validation on Modal L40S
pytest fastvideo/tests/train/utils/test_lora.py -v -k cuda

git diff --check origin/main...HEAD

Modal runs:

• Two-GPU CUDA/NCCL test: https://modal.com/apps/q447747035/main/ap-oiSFcZLlI9ho4hOgdBZRyK
• Full pre-commit run --all-files: https://modal.com/apps/q447747035/main/ap-qaOpzqleeamWskflENn5xG

Test Results

Test output

Local distributed LoRA tests:

fastvideo/tests/train/utils/test_lora.py::test_replicated_lora_parameter_stays_consistent_across_ranks PASSED
fastvideo/tests/train/utils/test_lora.py::test_replicated_lora_parameter_stays_consistent_on_cuda SKIPPED

1 passed, 1 skipped

Complete modular training utility suite:

76 passed, 1 skipped, 14 warnings in 11.37s

Modal L40S:2 CUDA/NCCL validation:

fastvideo/tests/train/utils/test_lora.py::test_replicated_lora_parameter_stays_consistent_on_cuda PASSED

1 passed, 1 deselected, 14 warnings in 18.39s

Full pre-commit run on Linux against commit 899383b9:

yapf                                      Passed
ruff                                      Passed
codespell                                 Passed
PyMarkdown                                Passed
Lint GitHub Actions workflow files        Passed
mypy                                      Passed
Check for spaces in all filenames         Passed
Suggestion                                Passed

git diff --check origin/main...HEAD

Passed

The warnings are existing PyTorch torch.jit.script_method deprecation warnings.

Checklist

• I ran pre-commit run --all-files and fixed all issues
• I added or updated tests for my changes
• I updated documentation if needed
  • No user-facing configuration or API changed, so documentation updates are not required.
• I considered GPU memory impact of my changes
  • The change adds no persistent GPU allocations. It introduces small gradient collectives for LoRA parameters.

For model/pipeline changes:

• SSIM regression tests are not applicable because model outputs and inference behavior are unchanged
• Support matrix updates are not applicable because no model support was added or changed

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]