Add V-JEPA 2 (Meta FAIR) distributed training test case

[paragao]

Summary

• Add V-JEPA 2 (Meta FAIR) ViT-g/16 1B-parameter self-supervised video model as a new PyTorch distributed training test case
• Includes Slurm (Pyxis/Enroot) and Kubernetes (PyTorchJob) deployment manifests
• Benchmarked on 8x p5en.48xlarge (64x NVIDIA H200 GPUs)

What is V-JEPA 2?

V-JEPA 2 is Meta FAIR's self-supervised video model that learns visual representations by predicting masked video patches. It achieves state-of-the-art on motion understanding and human action anticipation benchmarks. The ViT-g/16 variant has 1.03B encoder parameters.

Files Added

3.test_cases/pytorch/vjepa2/
├── vjepa2.Dockerfile                      # NVIDIA PyTorch 25.03 base (CUDA 13, Python 3.11)
├── README.md                              # Full walkthrough with benchmark results
├── slurm/
│   ├── benchmark_training.sbatch          # 200-iter benchmark (8 nodes)
│   ├── launch_training.sbatch             # Full 800-epoch pre-training
│   └── download_dataset.sbatch            # SSv2 dataset preparation
├── kubernetes/
│   └── vjepa2-benchmark.yaml              # PyTorchJob for EKS clusters
├── configs/
│   ├── benchmark-vitg-8nodes.yaml         # Quick benchmark config
│   └── pretrain-vitg-256px-16f.yaml       # Full pre-training config
└── scripts/
    ├── run_train.py                       # Thin srun-compatible launcher
    ├── generate_synthetic_dataset.py      # Synthetic video generator
    ├── prepare_ssv2.py                    # SSv2 CSV preparation
    ├── parse_benchmark.py                 # Log parser for throughput/MFU
    └── test_decord.py                     # Verify decord video loading

Key Technical Details

Launch pattern: V-JEPA 2 uses srun directly (not srun + torchrun). The run_train.py launcher calls app.vjepa.train.main() directly, which reads SLURM_LOCALID/SLURM_NTASKS/SLURM_PROCID for distributed setup. This avoids a bug in app/main.py where its subprocess launcher passes world_size=1 regardless of SLURM configuration.

Dataset: Supports both Something-Something v2 (SSv2) real data and synthetic generated videos for benchmarking.

Benchmark Results (8x p5en.48xlarge, 64x H200)

Metric	Value
Global batch size	1,536
Precision	BF16
Peak GPU memory	~32.9 GB / 143 GB

Testing

Validated on ParallelCluster with 8x p5en.48xlarge nodes running Slurm + Pyxis/Enroot with EFA networking. Job ran 200 iterations to completion with all 64 ranks correctly initialized via NCCL over EFA.

[KeitaW]

Review Batch 1/3 — Structure & Repository Hygiene

Thanks for this thorough contribution, Paulo! The utility scripts and READMEs are excellent quality. I have some structural and reproducibility findings below.

Significant code duplication between vjepa2/ and vjepa2.1/

These two directories share a large amount of identical code:

• scripts/generate_synthetic_dataset.py — identical (same git blob 74f922445)
• scripts/parse_benchmark.py — identical (same git blob 957b9efdf)
• scripts/prepare_ssv2.py — identical (same git blob 633288d17)
• scripts/test_decord.py — identical (same git blob 4881d1647)
• scripts/run_train.py — nearly identical (V-JEPA 2.1 adds 4 lines)
• Dockerfiles — nearly identical structure
• Slurm sbatch scripts — same structure, differing only in paths/config references

The repo convention says to "extend the existing test case — add platform-specific subdirectories, parameterize scripts for additional models, or add configuration variants — rather than creating a parallel directory tree with duplicated Dockerfiles, training scripts, and utilities."

I'd suggest consolidating into a single vjepa2/ directory that supports both V-JEPA 2 and 2.1 via different configs. The run_train.py launcher already dispatches based on the app field in the config (vjepa vs vjepa_2_1), so both versions can share the same launcher, scripts, Dockerfile, and sbatch templates. The V-JEPA 2.1 additions (image co-training, synthetic image generator) would simply add to the existing directory.

Missing license headers on README and config files

Both README.md files and all 4 configs/*.yaml files are missing license headers. The Slurm scripts, Python files, K8s manifests, and Dockerfiles all have them, so this is just an oversight. I'd suggest adding the standard header as a YAML comment in configs and HTML comment in READMEs.

[KeitaW]

Review Batch 2/3 — Deployment Pipeline

[KeitaW]

Review Batch 3/3 — Documentation Consistency

---

Things That Look Great

• Comprehensive utility scripts: The synthetic data generators (video and image), SSv2 CSV preparer, benchmark log parser, and decord test script form a complete toolkit that makes this test case truly self-contained.
• Excellent README documentation: Both READMEs walk through every step from dataset prep to result parsing, with clear architecture notes explaining the srun direct launch pattern and why app/main.py doesn't work with SLURM.
• Smart launch pattern: Using app.scaffold.main() to dispatch based on the config's app field is elegant and avoids the world_size=1 bug in app/main.py.
• Proper license headers on most files: Scripts, Dockerfiles, sbatch files, and K8s manifests all have the standard copyright header.
• HyperPod auto-resume detection: The if [ -d "/opt/sagemaker_cluster" ] pattern in sbatch scripts correctly detects HyperPod clusters and enables auto-resume.
• Both Slurm and Kubernetes deployment paths: Providing PyTorchJob manifests alongside Slurm scripts makes this accessible to EKS-based clusters too.
• Well-structured config separation: Benchmark configs (200 iterations, no checkpointing) vs. full pre-training configs (800+ epochs, regular checkpoints) give users clear starting points for different use cases.
• V-JEPA 2.1 comparison table: The feature comparison table in the V-JEPA 2.1 README clearly explains what changed between versions.

[KeitaW]

Few comments

[paragao]

@KeitaW please check that everything has been properly addressed so we can merge this.

[KeitaW]

Re-Review Batch 1/3 — Acknowledgments + Structure & Repository Hygiene

Thanks Paulo for working through the previous round of feedback. Most of the smaller items from the first pass are now resolved — pinned EFA, pinned vjepa2 commit, pinned pip packages, license headers added, eth removed from NCCL_SOCKET_IFNAME, awslabs URL fix, :latest removed from K8s images, set -euo pipefail adopted, and yaml.SafeLoader swapped in. Nice cleanup pass.

The main item still outstanding is the structural one: vjepa2/ and vjepa2.1/ are still two parallel trees with mostly-identical Dockerfiles, scripts, and sbatch files. Several new files added in the recent commits (FSDP launcher, nsys wrappers, B200 sbatches, sweep configs) duplicate that surface again. I'd really like to land that consolidation before merge — it's the single biggest maintenance lever for this contribution.

What's been addressed since the last review

• EFA pinned to 1.47.0 (meets CI minimum)
• vjepa2 git clone pinned to commit 204698b4...
• All pip packages pinned to tested versions
• License headers added to READMEs, configs, and .gitignore
• NCCL_SOCKET_IFNAME no longer excludes eth* (now ^docker,lo,veth everywhere)
• set -euo pipefail adopted across sbatch scripts
• K8s :latest replaced with :vjepa2 / :vjepa2.1 tags
• Stale aws-samples URL replaced with awslabs in both READMEs
• yaml.SafeLoader instead of yaml.FullLoader in launchers

---

Code duplication between vjepa2/ and vjepa2.1/ is still the headline issue

Looking at the current state, the duplication has actually grown since the last review. The two trees now share, in addition to the originally-flagged scripts, near-identical copies of:

• vjepa2.Dockerfile and vjepa2_1.Dockerfile — the only meaningful difference is one extra pinned package (Pillow==12.0.0); the vjepa2_1.Dockerfile even comments at the top "V-JEPA 2.1 reuses the same container as V-JEPA 2, since both training apps live in the same facebookresearch/vjepa2 repository"
• scripts/nsys_wrapper.sh — identical
• slurm/benchmark_training_b200.sbatch, slurm/nsys_profile.sbatch, slurm/nsys_profile_b200.sbatch, slurm/optim_sweep_b200.sbatch — same skeleton, only paths and config refs differ
• README sections on prerequisites, container build, parsing, and profiling

The vjepa2.1/README.md even reaches across the boundary several times — ## 2. Datasets says "See the V-JEPA 2 test case for SSv2 download" and the benchmark-results section points at ../vjepa2/benchmarks/vjepa2-benchmark-results.md. Cross-references like that are a strong tell that the directories want to be one.

run_train.py already dispatches via app.scaffold.main(params["app"], ...) based on the YAML's app field, so a single launcher (and Dockerfile, and sbatch templates) can serve both app: vjepa and app: vjepa_2_1 configs. I'd suggest restructuring as:

3.test_cases/pytorch/vjepa2/
├── README.md                    # covers both 2 and 2.1
├── vjepa2.Dockerfile            # single file, includes Pillow
├── configs/
│   ├── vjepa2/...               # app: vjepa configs
│   └── vjepa2.1/...             # app: vjepa_2_1 configs
├── scripts/
│   ├── run_train.py             # current scaffold-based launcher
│   ├── generate_synthetic_dataset.py
│   ├── generate_synthetic_images.py
│   ├── prepare_ssv2.py
│   ├── parse_benchmark.py
│   ├── nsys_wrapper.sh
│   └── test_decord.py
└── slurm/                       # one set of sbatches, parameterized via env vars

This is the same pattern called out in CONTRIBUTING.md ("extend the existing test case … rather than creating a parallel directory tree with duplicated Dockerfiles, training scripts, and utilities"). Happy to help if there's a structural reason this is harder than it looks.

Reconsider whether run_train_fsdp.py should ship at all

The FSDP path adds a 718-line bespoke training loop (vjepa2.1/scripts/run_train_fsdp.py), plus a config and sbatch script. The README itself says:

Benchmark finding: FSDP was found to be ~2x slower than baseline DDP in benchmarks (iter ~8,200 ms vs ~4,075 ms on video ranks). … The baseline DDP config is the recommended configuration for V-JEPA 2.1.

If the recommendation is "don't use this", I'd lean toward dropping the FSDP variant from the PR and reintroducing it later if/when it becomes useful. Carrying ~900+ lines of "for reference, but slower" code in awsome-distributed-training increases the maintenance surface (NCCL/PyTorch/EFA upgrades will need to keep it green) without users having a reason to run it. If you want to keep the finding documented, a short paragraph in the README is probably enough.

optim_sweep_b200.sbatch reads as developer experimentation, not a reference workflow

vjepa2.1/slurm/optim_sweep_b200.sbatch (+222 lines) and vjepa2/slurm/optim_sweep_b200.sbatch (+171 lines) iterate over phase configs and tweak env vars to sweep optimizations. They look like the harness you used to produce the benchmark numbers, not something a downstream user would run. The repo convention favors test cases that demonstrate a single, reproducible workflow; sweeps usually live in a personal branch or a README appendix.

I'd suggest either dropping these (along with the phase1/phase4 configs), or — if you want to keep the sweep methodology in-tree — moving it into a tools/ subdirectory with a README that frames it as an experimentation example, not a recommended config.

[KeitaW]

Re-Review Batch 2/3 — Deployment Pipeline (K8s / Slurm)

Three concrete issues introduced or surfaced by the most recent commits:

1. The new set -euo pipefail interacts badly with bare ${PYTHONPATH} in the B200 sbatches — set -u will cause the script to exit before srun if PYTHONPATH is unset on the user's environment (the common case). Fix is the standard ${PYTHONPATH:-} guard, suggested inline on each affected file.
2. The apt install libnccl-dev … || true line in both Dockerfiles silently swallows failures and doesn't pin a version, while CI requires NCCL >= 2.28 and the base nvcr.io/nvidia/pytorch:25.03-py3 ships NCCL 2.25 (per your own commit message c0fca12). Worth either dropping the || true so failures surface, or making it explicit that NCCL is provided by the base image.
3. Minor: the K8s image tags vjepa2:vjepa2 / vjepa2:vjepa2.1 use the variant name as the tag, which is unusual. A version-flavored tag (e.g. mirroring the upstream commit pin or the base image date) sets a better template for users.

[KeitaW]

Re-Review Batch 3/3 — Documentation Consistency + Things That Look Great

One cross-cutting documentation item, plus inline picks below:

B200 setup is documented only inside the sbatch comments

The B200 sbatches (vjepa2/slurm/benchmark_training_b200.sbatch, _optimized.sbatch, vjepa2.1/slurm/benchmark_training_b200.sbatch, _fsdp.sbatch) require a separate NeMo container, a cloned vjepa2_code/ directory on FSx, and a pip install --target /fsx/.../vjepa_deps step. None of this is mentioned in either README, so a user going directly through Build Container → Run Benchmark with B200 instances will hit unexplained CONTAINER_IMAGE paths. I'd suggest adding a short B200 setup subsection to the README, or at least pointing to the sbatch's header comment from the README.

---

Things That Look Great

• Genuine engagement with the previous review — all the small/medium items are visibly fixed, with commit messages that explain the why (e.g. 9a946379 calls out the cluster container freeze as the source for pinned versions).
• The B200 vs H200 split is well-motivated — keeping the H200 Dockerfile path simple while running B200 against a NeMo container with a known-compatible NCCL/EFA stack is a pragmatic choice and the commit message c0fca12 explains it clearly.
• Optimization knobs in run_train.py are tasteful — env-var-gated monkey-patches (VJEPA_TF32, VJEPA_FUSED_OPTIMIZER, VJEPA_GRAD_BUCKET_VIEW, VJEPA_PREFETCH_FACTOR) keep the upstream code unmodified while making perf experiments easy to compose.
• GradScaler-for-BF16 fix with the subclass approach (so Apex's GradScaler still inherits cleanly) is a nice touch and documented inline.
• .gitignore for benchmarks/ and profiling/ keeps generated artifacts out of the repo while still letting you keep the directory referenced from the README.
• Honest documentation of negative results — calling out that FSDP was 2× slower and torch.compile + activation checkpointing was 55% slower is exactly the kind of empirical guidance that makes this repo useful (even if I think the FSDP code itself probably doesn't need to ship — see Batch 1).
• Synthetic-dataset sizing guidance in the README (50K minimum, with the rationale about dataloader re-init) is the kind of subtle benchmarking pitfall that's exactly worth documenting.

---

Suggested priority before merge

1. Consolidate vjepa2/ and vjepa2.1/ into a single test case (Batch 1) — the headline blocker, and the rationale for several other findings.
2. Fix the set -u × ${PYTHONPATH} bug in the four B200 sbatches (Batch 2) — small but currently breaks the scripts on common environments.
3. Decide what to do with run_train_fsdp.py and optim_sweep_b200.sbatch (Batch 1) — drop or relocate.
4. Address the smaller documentation/Dockerfile items at your discretion.

[paragao]

@KeitaW addressed comments. I decided to maintain both vjepa2 and vjepa2.1 folders, as this is how they are tracked on the original forked repo. Those are two different approaches to world models, where vjepa2 only uses images and vjepa2.1 combines videos and images. All duplication removed.