Adding DDP to the repo for multi-GPU training

[vratins]

Multi-GPU flow matching training via torchrun

New src/distributed.py (shared DDP helpers) + DDP wiring in train.py. No CLI flag — torchrun env vars are the only switch, single-GPU runs unchanged. Verified on 2×H100. tests are CPU-only for CI for now.

Updated README.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: bb8c8452-564f-456f-9c22-eb83658c0643

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

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.}

[Copilot]

Pull request overview

Adds first-class multi-GPU (torchrun/DDP) support to the training pipeline by introducing shared distributed utilities, wiring DDP into scripts/train.py, and extending data loading + metrics reduction so multi-process runs behave like single-process runs from a user/config perspective.

Changes:

• Introduces src/distributed.py helpers for launcher detection, barriers, single-writer work, and mean-reduction of metrics.
• Updates scripts/train.py to support torchrun-based DDP (sampler sharding, reduced metrics, controlled logging/checkpointing, and gradient sync scheduling) plus a rank-0 cache prebuild phase.
• Adds CPU-only tests for distributed helper logic and cache sharding, and documents DDP usage in the README.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/test_distributed.py	Adds unit tests for DDP helper behavior, sampler wiring, grad-sync scheduling, and cache sharding logic.
src/flow.py	Fixes cutoff lookup to work correctly when the model is wrapped by DDP.
src/distributed.py	New module: DDP activation detection, process-group setup/teardown, CPU store “run once” helper, and all-reduce mean reduction.
src/dataset.py	Extends get_dataloader with optional sampler + distributed=True to build a DistributedSampler automatically.
scripts/train.py	Wires in DDP end-to-end (DDP wrapping, reduced metrics, per-rank IO/logging rules, distributed samplers, cache prebuild, and grad sync control).
README.md	Documents how to run multi-GPU training with torchrun and what changes under DDP.
pyproject.toml	Adds a ty override for src/distributed.py to avoid false positives from torch.distributed typing.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.