From 2cb832290a971058ce02546205eeebdea69f3433 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 20 Jul 2026 12:28:34 +0300 Subject: [PATCH] [bugfix]: fit Kandinsky5 e2e stage 2 in one 80GiB GPU; make multi-GPU sharding viable The generator-update fix costs a second full Adam state (~16GiB for the 2B student) on top of the critic's: a recorded stage-2 run peaked at ~76GiB and OOMed allocating the last MiBs with ~1.9GiB lost to fragmentation. Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True for the training stages to reclaim that headroom. Also make KANDINSKY5_E2E_NUM_GPUS=2+ actually work as the fallback: DP_SP_BatchSampler floor-divides batch count by the number of DP groups with drop_last=True, so the previous 1-row dataset produced an EMPTY dataloader on every rank at 2+ GPUs. Synthesize max(2, NUM_GPUS) manifest rows pointing at the same clip -- identical latents/caption keep the single-sample-overfit semantics. Co-Authored-By: Claude Fable 5 --- .../test_e2e_kandinsky5_dmd_t2v_overfit.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py index 5399b7f34..e1cf04491 100644 --- a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py +++ b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py @@ -198,11 +198,20 @@ def _synthesize_single_sample() -> None: writer.write(frame) writer.release() + # One physical clip, but max(2, NUM_GPUS) manifest rows all pointing at + # it: DP_SP_BatchSampler (fastvideo/dataset/parquet_dataset_map_style.py) + # floor-divides the number of batches by the number of data-parallel + # groups with drop_last=True, so a dataset with fewer rows than GPUs + # yields an EMPTY dataloader on every rank. Duplicate rows keep the + # single-sample-overfit semantics (identical latents/caption) while + # making KANDINSKY5_E2E_NUM_GPUS=2+ a working escape hatch when one GPU + # doesn't have the memory headroom for stage 2 (see _run_stage). + num_rows = max(2, int(NUM_GPUS)) with open(RAW_DATA_DIR / "videos2caption.json", "w") as f: json.dump([{ "path": "sample_0.mp4", "cap": [GENERATION_PROMPT], - }], f) + }] * num_rows, f) def _run_preprocessing() -> None: @@ -275,6 +284,14 @@ def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_ov *extra_overrides, ] env = dict(os.environ) + # Stage 2 holds three 2B-param models plus TWO full Adam states (the + # student's only exists because generator_update_interval is overridden + # to 1 above -- a recorded run peaked at ~76 GiB and died allocating + # the last MiBs on an 80 GiB device with ~1.9 GiB lost to + # fragmentation). Expandable segments reclaims that headroom; if a + # single device still can't fit, shard with KANDINSKY5_E2E_NUM_GPUS=2+ + # (the synthesized manifest guarantees >= one row per rank). + env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") env.update(env_overrides) subprocess.run(cmd, cwd=str(REPO_ROOT), env=env, check=True)