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]

This comment is the strongest argument for consolidating the two directories

If 2.1 is acknowledged in-tree as reusing the V-JEPA 2 container, having a separate vjepa2_1.Dockerfile (whose only delta from vjepa2.Dockerfile is one extra pinned package) creates ongoing drift risk — a CUDA/EFA/NCCL bump applied to one will need to be repeated on the other. I'd suggest deleting this file in favor of a single Dockerfile that includes Pillow==12.0.0 (it's harmless on the V-JEPA 2 side too).

See the structural finding in the review body for the proposed unified layout.

[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]

Bare ${PYTHONPATH} will break under set -u

With the new set -euo pipefail at the top of the script, this line will fail with PYTHONPATH: unbound variable on any node where PYTHONPATH isn't already set in the user's environment — which is the common case under Slurm/Pyxis. The script will exit before srun even runs.

Use the standard guard:

Suggested change

	export PYTHONPATH=/workspace:${VJEPA_DEPS}:${VJEPA2_CODE}:${PYTHONPATH}
	export PYTHONPATH=/workspace:${VJEPA_DEPS}:${VJEPA2_CODE}:${PYTHONPATH:-}

A quick env -i bash -c 'set -u; . path/to/script.sbatch' smoke test should catch the rest of these.

[KeitaW]

Same set -u × bare ${PYTHONPATH} issue as benchmark_training_b200.sbatch

Suggested change

	export PYTHONPATH=/workspace:${VJEPA_DEPS}:${VJEPA2_CODE}:${PYTHONPATH}
	export PYTHONPATH=/workspace:${VJEPA_DEPS}:${VJEPA2_CODE}:${PYTHONPATH:-}

[KeitaW]

Same set -u × bare ${PYTHONPATH} issue

Suggested change

	export PYTHONPATH=${VJEPA_DEPS}:${VJEPA2_CODE}:${PYTHONPATH}
	export PYTHONPATH=${VJEPA_DEPS}:${VJEPA2_CODE}:${PYTHONPATH:-}

[KeitaW]

Same set -u × bare ${PYTHONPATH} issue

Suggested change

	export PYTHONPATH=${VJEPA_DEPS}:${VJEPA2_CODE}:${PYTHONPATH}
	export PYTHONPATH=${VJEPA_DEPS}:${VJEPA2_CODE}:${PYTHONPATH:-}

[KeitaW]

Silent-fail apt install libnccl-dev with no NCCL version pin

Two concerns:

1. The trailing || true masks any apt failure — if libnccl-dev is missing from the index or fails to install, the build still succeeds and the user only finds out at training time.
2. There's no explicit NCCL version pin, while the CI version check requires NCCL >= 2.28. The base nvcr.io/nvidia/pytorch:25.03-py3 ships NCCL 2.25 (per commit message c0fca12), and apt install libnccl-dev from the default Ubuntu repo typically pulls a separately-versioned package.

I'd suggest dropping the || true so failures surface, or adding an explicit comment that NCCL is provided by the base image and the apt step is best-effort. Either way, document which NCCL version the final image actually carries so users (and the CI version-check job) know the H200 reference path meets the minimum. The B200 path side-steps this by using a NeMo container, but the H200 reference path needs to demonstrably be >= 2.28.

[KeitaW]

Same silent-fail libnccl-dev install + no NCCL version pin as vjepa2.Dockerfile

See the comment on vjepa2/vjepa2.Dockerfile:23 for the full reasoning. (Another data point that these two Dockerfiles want to be one file.)

[KeitaW]

vjepa2:vjepa2 is an awkward tag convention

Using the variant name as the tag is unusual — most readers will expect the variant to be the repository name, with the tag carrying a version or commit. I'd suggest aligning with the upstream commit you've pinned in the Dockerfile, e.g.:

Suggested change

	image: <ACCOUNT_ID>.dkr.ecr.<REGION>.amazonaws.com/vjepa2:vjepa2
	image: <ACCOUNT_ID>.dkr.ecr.<REGION>.amazonaws.com/vjepa2:204698b4

…or mirroring the base image's date tag (vjepa2:25.03-py3). Minor, but it sets a better template for users copy-pasting this into their own deployments.

[KeitaW]

Same awkward vjepa2:vjepa2.1 tag convention as vjepa2-benchmark.yaml

See the comment on vjepa2/kubernetes/vjepa2-benchmark.yaml:42. A version-flavored tag (commit SHA or date) reads more naturally than using the variant name as the tag.

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

[KeitaW]

Per-GPU video batch size is reported two different ways in the same README

This section says ranks 32-63 use batch size 24 per GPU. The Image/Video co-training rank split section further down (line 199) instead says batch size 48/GPU = 1,536 videos/step for the same ranks.

I think the difference is that the first is the YAML-configured batch_size and the second is the effective size after the V-JEPA 2.1 code auto-doubles it to keep the global batch constant when rank_ratio=0.5. That's worth saying once, in one place — right now the two numbers look like a contradiction. Either fold the auto-adjust sentence (line 201) up here, or drop the (batch size 24 per GPU) parenthetical.

[KeitaW]

Reference to a gitignored file is dead

benchmarks/vjepa2-benchmark-results.md is gitignored (per the .gitignore added in this PR), so it's never present in a fresh clone — readers following this line will hit nothing. If results are published elsewhere (a blog post, an internal doc, or the PR description's table), I'd link there instead. Otherwise it's cleaner to just drop the line.

Suggested change

_Benchmark results are maintained separately in `benchmarks/vjepa2-benchmark-results.md` (gitignored)._

[KeitaW]

Same dead reference to a gitignored file as vjepa2/README.md:284

The ../vjepa2/benchmarks/ directory is gitignored, so the file isn't present in a fresh clone. Drop the line, or replace it with a link to wherever the published benchmark results actually live.

Suggested change

_Benchmark results are maintained separately in `../vjepa2/benchmarks/vjepa2-benchmark-results.md` (gitignored)._