diff --git a/examples/distill/MiniMax-H3/overfit_vsa0.sh b/examples/distill/MiniMax-H3/overfit_vsa0.sh new file mode 100644 index 000000000..695a30885 --- /dev/null +++ b/examples/distill/MiniMax-H3/overfit_vsa0.sh @@ -0,0 +1,49 @@ +#!/bin/bash +# Single-sample VSA-H3 DMD2 overfit for MiniMax H3 (sparsity 0.0, parity vs +# dense). Two steps: +# (a) preprocess ONE mp4 (with audio) + prompt into the t2va parquet row +# the H3 dataloader consumes (768x1344, 124 frames @ 24 fps, 32 kHz); +# (b) train on 4 GPUs: student on VIDEO_SPARSE_ATTN_H3, teacher/critic on +# FLASH_ATTN with the FA4 path (FASTVIDEO_FA4=1). +# +# Usage: +# VIDEO=/path/clip.mp4 PROMPT="a caption" bash examples/distill/MiniMax-H3/overfit_vsa0.sh +# SKIP_PREPROCESS=1 bash examples/distill/MiniMax-H3/overfit_vsa0.sh # data already written + +set -euo pipefail + +export TOKENIZERS_PARALLELISM=false +export MASTER_PORT=${MASTER_PORT:-29513} + +NUM_GPUS=${NUM_GPUS:-4} +CONFIG=${CONFIG:-examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml} +DATA_DIR=${DATA_DIR:-/mnt/h3-dmd2-overfit/data} +OUTPUT_DIR=${OUTPUT_DIR:-outputs/minimax_h3_dmd2_vsa0_overfit} +MODEL_PATH=${MODEL_PATH:-data/models/MiniMax-H3} + +# ---------------------------------------------------------------- preprocess +if [ -z "${SKIP_PREPROCESS:-}" ]; then + VIDEO=${VIDEO:?set VIDEO=/path/to/clip.mp4 (with soundtrack) or SKIP_PREPROCESS=1} + PROMPT=${PROMPT:?set PROMPT="text prompt for the clip"} + python -m fastvideo.pipelines.preprocess.preprocess_minimax_h3_overfit \ + --video "$VIDEO" \ + --prompt "$PROMPT" \ + --output-dir "$DATA_DIR" \ + --model-path "$MODEL_PATH" +fi + +# --------------------------------------------------------------------- train +# FA4 flash-attention path for the dense teacher/critic FLASH_ATTN roles. +export FASTVIDEO_FA4=1 + +torchrun \ + --nnodes 1 \ + --master_port "$MASTER_PORT" \ + --nproc_per_node "$NUM_GPUS" \ + -m fastvideo.train.entrypoint.train \ + --config "$CONFIG" \ + --training.distributed.num_gpus "$NUM_GPUS" \ + --training.distributed.sp_size "$NUM_GPUS" \ + --training.distributed.hsdp_shard_dim "$NUM_GPUS" \ + --training.data.data_path "$DATA_DIR" \ + --training.checkpoint.output_dir "$OUTPUT_DIR" diff --git a/examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml b/examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml new file mode 100644 index 000000000..fb907d8a7 --- /dev/null +++ b/examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml @@ -0,0 +1,125 @@ +# Single-sample DMD2 overfit: MiniMax H3 with the VSA-H3 student backend at +# sparsity 0.0 (parity-vs-dense experiment; the backend's sparsity-0 path is +# parity-proven <=2e-4 vs dense). +# +# - Student: VIDEO_SPARSE_ATTN_H3 backend; DMD2 forwards it attn_kind="vsa" +# with real packed-sequence metadata built in MiniMaxH3DMDModel.prepare_batch. +# - Teacher/critic: FLASH_ATTN (run with FASTVIDEO_FA4=1 for the FA4 path); +# their forwards stay dense (attn_metadata=None). +# - Data: ONE preprocessed t2va parquet row (768x1344, 124 frames @ 24 fps) +# produced by fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py +# — see examples/distill/MiniMax-H3/overfit_vsa0.sh. + +models: + student: + _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel + init_from: MiniMaxAI/MiniMax-H3 + trainable: true + enable_gradient_checkpointing_type: full + attention_backend: VIDEO_SPARSE_ATTN_H3 + teacher: + _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel + init_from: MiniMaxAI/MiniMax-H3 + trainable: false + disable_custom_init_weights: true + attention_backend: FLASH_ATTN + critic: + _target_: fastvideo.train.models.minimax_h3.MiniMaxH3DMDModel + init_from: MiniMaxAI/MiniMax-H3 + trainable: true + disable_custom_init_weights: true + enable_gradient_checkpointing_type: full + attention_backend: FLASH_ATTN + +method: + _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method + rollout_mode: data_latent + generator_update_interval: 5 + real_score_guidance_scale: 3.5 + dmd_denoising_steps: [1000, 757, 522] + min_timestep_ratio: 0.02 + max_timestep_ratio: 0.98 + # H3 has no negative-prompt encoder at training time; unconditional teacher + # forwards zero the text embeddings instead. + cfg_uncond: + text: zero + + # Critic optimizer (required — no fallback to training.optimizer) + fake_score_learning_rate: 8.0e-6 + fake_score_betas: [0.0, 0.999] + fake_score_lr_scheduler: constant + +training: + distributed: + num_gpus: 4 + sp_size: 4 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + # Preprocessed-latents parquet dir; the cluster run produces it. + data_path: /mnt/h3-dmd2-overfit/data + preprocessed_data_type: t2va + dataloader_num_workers: 0 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 42 + num_latent_t: 37 + num_height: 768 + num_width: 1344 + num_frames: 124 + + # Student VSA sparsity (0.0 = every tile selected, dense parity). + vsa: + sparsity: 0.0 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.0, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 2000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/minimax_h3_dmd2_vsa0_overfit + training_state_checkpointing_steps: 250 + checkpoints_total_limit: 3 + + tracker: + project_name: h3-dmd2-vsa + run_name: dmd2_vsa0_overfit + + model: + precondition_outputs: false + enable_gradient_checkpointing_type: full + + dit_precision: bf16 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + validation: + _target_: fastvideo.train.callbacks.validation.ValidationCallback + pipeline_target: fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.MiniMaxH3Pipeline + dataset_file: examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json + every_steps: 250 + sampling_steps: [3] + guidance_scale: 1.0 + num_frames: 124 + num_videos_per_prompt: 1 + use_validation_media_conditioning: false + offload_training_state: true + text_encoder_cpu_offload: true + vae_cpu_offload: true + ema: + _target_: fastvideo.train.callbacks.ema.EMACallback + decay: 0.98 + start_iter: 0 + +pipeline: {} diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index cdb75840a..2bd73d180 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -261,6 +261,15 @@ class MiniMaxH3VSAImpl(AttentionImpl): n_tiles = attn_metadata.variable_block_sizes.numel() target_shape = (x.shape[0], n_tiles * _TILE_ELEMS, x.shape[-2], x.shape[-1]) + if torch.is_grad_enabled() and x.requires_grad: + # Grad-tracking forwards (VSA-H3 training students) get a fresh + # buffer per call: autograd saves views of the tiled buffer for + # backward (layer.py chunks q/k/v out of it), so a later layer's + # in-place overwrite of a shared buffer raises the saved-tensor + # version error at backward. Mirrors Wan training's + # vsa_cache_tile_buf=False default. + return scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, None) + # single scatter: untile_combined_index maps original row i to its # padded slot, so this is exactly the inverse of postprocess_output holder = attn_metadata.tile_buf_holder diff --git a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py index 3a8f706da..655e99329 100644 --- a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py @@ -313,21 +313,35 @@ def validate_preprocessed_training_data( print(f"Validated one Crush-Smol H3 training row in {expected_path}") -def main() -> None: - """Encode the selected Crush-Smol record with each H3 component in sequence. +def main( + video: Path | None = None, + prompt: str | None = None, + output_dir: Path = OUTPUT_DIR, + model_path: Path = MODEL_PATH, +) -> None: + """Encode one audio-video-caption record with each H3 component in sequence. + With no arguments this encodes the pinned Crush-Smol record; ``--video`` + + ``--prompt`` preprocess an arbitrary mp4 (with soundtrack) instead. Releasing each component before loading the next component keeps video, audio, and text preprocessing within one GPU's memory. """ _init_single_process_distributed() - resolved_model_path = MODEL_PATH.resolve() + resolved_model_path = model_path.resolve() if not resolved_model_path.is_dir(): raise FileNotFoundError(f"Filtered MiniMax H3 model directory is missing at {resolved_model_path}") model_index = verify_model_config_and_directory(str(resolved_model_path)) - video_path, caption = load_crush_smol_training_sample( - DATA_DIR / "videos2caption.json", - DATA_DIR / "videos", - ) + if video is None: + video_path, caption = load_crush_smol_training_sample( + DATA_DIR / "videos2caption.json", + DATA_DIR / "videos", + ) + else: + if not prompt or not prompt.strip(): + raise ValueError("--prompt is required when --video is given") + video_path, caption = video, prompt.strip() + if not video_path.is_file(): + raise FileNotFoundError(f"Training video is missing at {video_path}") frames, waveform = load_training_media(video_path) pipeline_config = MiniMaxH3PipelineConfig() fastvideo_args = FastVideoArgs( @@ -345,21 +359,34 @@ def main() -> None: audio_latents = encode_audio_latents(waveform, resolved_model_path, model_index, fastvideo_args) text_embedding = encode_text_embedding(caption, resolved_model_path, model_index, fastvideo_args) record = build_parquet_record( - file_name=TRAINING_VIDEO_NAME, + file_name=video_path.name, caption=caption, video_latents=video_latents, audio_latents=audio_latents, text_embedding=text_embedding, ) - output_path = write_parquet(record, OUTPUT_DIR) - print(f"Wrote one Crush-Smol MiniMax H3 T2VA record to {output_path}") + output_path = write_parquet(record, output_dir) + print(f"Wrote one MiniMax H3 T2VA record to {output_path}") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--validate-only", action="store_true") + parser.add_argument("--video", + type=Path, + default=None, + help="mp4 with a soundtrack (>=124 frames at 24 fps); " + "default: the pinned Crush-Smol record") + parser.add_argument("--prompt", type=str, default=None, help="caption for --video") + parser.add_argument("--output-dir", type=Path, default=OUTPUT_DIR) + parser.add_argument("--model-path", type=Path, default=MODEL_PATH) cli_args = parser.parse_args() if cli_args.validate_only: validate_preprocessed_training_data() else: - main() + main( + video=cli_args.video, + prompt=cli_args.prompt, + output_dir=cli_args.output_dir, + model_path=cli_args.model_path, + ) diff --git a/fastvideo/tests/attention/test_vsa_h3_metadata.py b/fastvideo/tests/attention/test_vsa_h3_metadata.py index b4b3efe5a..2fc1fe78f 100644 --- a/fastvideo/tests/attention/test_vsa_h3_metadata.py +++ b/fastvideo/tests/attention/test_vsa_h3_metadata.py @@ -133,9 +133,33 @@ def test_prefix_queries_stay_dense_at_high_sparsity(): "video rows should actually be sparse at 75%" +def test_tile_under_grad_does_not_reuse_shared_buffer(): + """Training forwards need fresh tile buffers; in-place reuse of the shared + holder would trip autograd's saved-tensor version check at backward.""" + meta = _build(_TINY) + impl = _impl() + seq = meta.total_seq_length + + x1 = torch.randn(1, seq, 2, 8, requires_grad=True) + x2 = torch.randn(1, seq, 2, 8, requires_grad=True) + buf1 = impl.tile(x1, meta) + saved = (buf1 * buf1).sum() # saves buf1 for backward, like the kernel + buf2 = impl.tile(x2, meta) + assert buf1 is not buf2 + saved.backward() # raises "modified by an inplace operation" on reuse + assert x1.grad is not None and x2.grad is None + + # no-grad paths (inference / rollout) keep the single-buffer reuse + with torch.no_grad(): + y1 = impl.tile(x1.detach(), meta) + y2 = impl.tile(x2.detach(), meta) + assert y1 is y2 + + if __name__ == "__main__": test_geometry_720p() test_mask_policy() test_sparsity_zero_matches_dense_sdpa() test_prefix_queries_stay_dense_at_high_sparsity() + test_tile_under_grad_does_not_reuse_shared_buffer() print("all VSA-H3 CPU checks passed") diff --git a/fastvideo/tests/train/methods/test_minimax_h3_dmd2.py b/fastvideo/tests/train/methods/test_minimax_h3_dmd2.py index 9931664ef..1a854a266 100644 --- a/fastvideo/tests/train/methods/test_minimax_h3_dmd2.py +++ b/fastvideo/tests/train/methods/test_minimax_h3_dmd2.py @@ -14,6 +14,9 @@ import pytest import torch import yaml +from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadata +from fastvideo.forward_context import get_forward_context +from fastvideo.platforms import AttentionBackendEnum from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method from fastvideo.train.models.minimax_h3 import MiniMaxH3DMDModel, MiniMaxH3Model from fastvideo.train.models.minimax_h3.minimax_h3 import shift_noise_amount @@ -22,6 +25,7 @@ from fastvideo.train.utils.config import load_run_config _FIXTURE = Path(__file__).resolve().parent.parent / "fixtures" / "minimax_h3_dmd2_min.yaml" _REPO_ROOT = Path(__file__).resolve().parents[4] _EXPERIMENT_CONFIG = _REPO_ROOT / "examples/train/configs/distribution_matching/minimax_h3/dmd2_t2va.yaml" +_VSA_OVERFIT_CONFIG = (_REPO_ROOT / "examples/train/configs/distribution_matching/minimax_h3/dmd2_vsa0_overfit.yaml") # Fixture geometry: video latents [1, 24, 2, 4, 4] and audio latents # [1, 2, 32, 8]; the packed adapter stores video-major [1, T, C, H, W]. @@ -39,9 +43,11 @@ class _TinyJointTransformer(torch.nn.Module): super().__init__() self.scale = torch.nn.Parameter(torch.tensor(scale)) self.last_encoder_hidden_states: torch.Tensor | None = None + self.last_attn_metadata = None def forward(self, **kwargs): self.last_encoder_hidden_states = kwargs["encoder_hidden_states"] + self.last_attn_metadata = get_forward_context().attn_metadata return ( kwargs["hidden_states"] * self.scale, kwargs["audio_hidden_states"] * self.scale, @@ -61,6 +67,7 @@ def _make_model( model.transformer = _TinyJointTransformer(scale) model.training_config = training_config model.sp_group = None + model.attention_backend = None return model @@ -73,6 +80,7 @@ def _tiny_training_config(): num_width=64, ), distributed=SimpleNamespace(sp_size=1), + vsa_sparsity=0.0, ) @@ -297,11 +305,131 @@ def test_uncond_forward_zeroes_text_and_guards_policies(monkeypatch: pytest.Monk model.set_requires_negative_conditioning(False) +# ---------------------------------------------------------------------- +# VSA-H3 wiring +# ---------------------------------------------------------------------- + + +def test_prepare_batch_builds_vsa_h3_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + """The VSA-H3 role gets real packed-sequence metadata; dense view stays None.""" + tc = _tiny_training_config() + tc.vsa_sparsity = 0.35 + model = _make_model(monkeypatch, tc) + model.attention_backend = AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3 + + batch = model.prepare_batch(_raw_batch(), generator=torch.Generator().manual_seed(7)) + + meta = batch.attn_metadata_vsa + assert isinstance(meta, MiniMaxH3VSAMetadata) + assert batch.attn_metadata is None + assert meta.VSA_sparsity == pytest.approx(0.35) + # Packed layout: 2 text rows | 0 condition rows | 16 stereo audio rows | + # 8 video rows ([1, 24, 2, 4, 4] latents at patch (1, 2, 2)). + assert meta.total_seq_length == 26 + assert meta.num_prefix_tiles == 2 + assert meta.num_video_tiles == 1 + assert meta.variable_block_sizes.tolist() == [2, 16, 8] + assert int(meta.variable_block_sizes.sum()) == meta.total_seq_length + + +def test_predict_noise_routes_vsa_metadata_by_attn_kind(monkeypatch: pytest.MonkeyPatch) -> None: + """Student "vsa" forwards see the VSA metadata; "dense" forwards see None.""" + model = _make_model(monkeypatch, _tiny_training_config()) + model.attention_backend = AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3 + batch = model.prepare_batch(_raw_batch(), generator=torch.Generator().manual_seed(7)) + noisy = torch.randn(1, _PACKED_NUMEL).to(torch.bfloat16) + timestep = torch.tensor([757], dtype=torch.long) + + model.predict_noise(noisy, timestep, batch, conditional=True, attn_kind="vsa") + assert model.transformer.last_attn_metadata is batch.attn_metadata_vsa + assert isinstance(model.transformer.last_attn_metadata, MiniMaxH3VSAMetadata) + + model.predict_noise(noisy, timestep, batch, conditional=True, attn_kind="dense") + assert model.transformer.last_attn_metadata is None + + +def _ctor_training_config() -> SimpleNamespace: + """The minimum surface MiniMaxH3Model.__init__ reads from TrainingConfig.""" + return SimpleNamespace( + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(uniform_parameter_dtype=False)), + data=SimpleNamespace( + train_batch_size=1, + training_cfg_rate=0.0, + preprocessed_data_type="t2va", + ), + model=SimpleNamespace(enable_gradient_checkpointing_type=None), + ) + + +def test_per_role_attention_backend_override_resolves(monkeypatch: pytest.MonkeyPatch) -> None: + """Each role's backend reaches the loader; unsupported backends fail fast.""" + captured: dict[str, AttentionBackendEnum | None] = {} + + def _fake_load(**kwargs): + captured[kwargs["model_path"]] = kwargs["attention_backend"] + return _TinyJointTransformer() + + monkeypatch.setattr( + "fastvideo.train.models.minimax_h3.minimax_h3.load_module_from_path", + _fake_load, + ) + + student = MiniMaxH3DMDModel( + init_from="role/student", + training_config=_ctor_training_config(), + trainable=True, + attention_backend="VIDEO_SPARSE_ATTN_H3", + ) + teacher = MiniMaxH3DMDModel( + init_from="role/teacher", + training_config=_ctor_training_config(), + trainable=False, + attention_backend="FLASH_ATTN", + ) + + assert student.attention_backend is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3 + assert teacher.attention_backend is AttentionBackendEnum.FLASH_ATTN + # load_module_from_path turns this request into the construction scope + # that binds the backend to the transformer's attention layers. + assert captured["role/student"] is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3 + assert captured["role/teacher"] is AttentionBackendEnum.FLASH_ATTN + assert not any(p.requires_grad for p in teacher.transformer.parameters()) + + with pytest.raises(ValueError, match="supports the attention backends"): + MiniMaxH3DMDModel( + init_from="role/bad", + training_config=_ctor_training_config(), + attention_backend="VIDEO_SPARSE_ATTN", + ) + + # ---------------------------------------------------------------------- # Config contracts # ---------------------------------------------------------------------- +def test_vsa0_overfit_config_pins_roles_and_experiment() -> None: + """The overfit config runs student VSA-H3 at sparsity 0, dense FA roles.""" + config = yaml.safe_load(_VSA_OVERFIT_CONFIG.read_text()) + models, method, training = config["models"], config["method"], config["training"] + + assert models["student"]["attention_backend"] == "VIDEO_SPARSE_ATTN_H3" + assert models["teacher"]["attention_backend"] == "FLASH_ATTN" + assert models["critic"]["attention_backend"] == "FLASH_ATTN" + assert models["teacher"]["trainable"] is False + assert training["vsa"]["sparsity"] == 0.0 + assert method["rollout_mode"] == "data_latent" + assert method["generator_update_interval"] == 5 + assert method["dmd_denoising_steps"] == [1000, 757, 522] + assert training["data"]["train_batch_size"] == 1 + assert training["data"]["data_path"] == "/mnt/h3-dmd2-overfit/data" + assert training["data"]["num_height"] == 768 + assert training["data"]["num_width"] == 1344 + assert training["data"]["num_frames"] == 124 + assert training["loop"]["max_train_steps"] == 2000 + assert training["tracker"]["project_name"] == "h3-dmd2-vsa" + + def test_h3_dmd2_fixture_resolves_trio_contract() -> None: """The fixture must wire the H3 DMD trio through the modular builder path.""" config = load_run_config(str(_FIXTURE)) diff --git a/fastvideo/tests/train/utils/test_tracking.py b/fastvideo/tests/train/utils/test_tracking.py new file mode 100644 index 000000000..9f2053815 --- /dev/null +++ b/fastvideo/tests/train/utils/test_tracking.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for the modular trainer's wandb tracking guard. + +The Trainer logs each step's loss dict through ``build_tracker``'s tracker on +rank 0; these tests pin the graceful no-op path (wandb missing / no +credentials) and the logged loss-dict shape with a monkeypatched wandb. +""" + +import builtins +import sys +from types import SimpleNamespace + +from fastvideo.train.utils import tracking +from fastvideo.train.utils.training_config import ( + CheckpointConfig, + TrackerConfig, +) +from fastvideo.training.trackers import DummyTracker, WandbTracker + + +class _FakeRun: + + def __init__(self, **init_kwargs): + self.init_kwargs = init_kwargs + self.logged = [] + + def log(self, metrics, step=None): + self.logged.append((metrics, step)) + + def finish(self): + pass + + +def _fake_wandb(): + runs = [] + + def _init(**kwargs): + run = _FakeRun(**kwargs) + runs.append(run) + return run + + return SimpleNamespace(init=_init, runs=runs, api=SimpleNamespace(api_key="key")) + + +def test_wandb_usable_false_when_not_importable(monkeypatch) -> None: + real_import = builtins.__import__ + + def _blocked(name, *args, **kwargs): + if name == "wandb": + raise ImportError("no wandb") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _blocked) + monkeypatch.delitem(sys.modules, "wandb", raising=False) + assert tracking._wandb_usable() is False + + +def test_wandb_usable_false_without_credentials(monkeypatch) -> None: + monkeypatch.delenv("WANDB_API_KEY", raising=False) + monkeypatch.delenv("WANDB_MODE", raising=False) + monkeypatch.setitem( + sys.modules, + "wandb", + SimpleNamespace(api=SimpleNamespace(api_key=None)), + ) + assert tracking._wandb_usable() is False + + +def test_wandb_usable_with_env_key(monkeypatch) -> None: + monkeypatch.setenv("WANDB_API_KEY", "key") + monkeypatch.setitem(sys.modules, "wandb", SimpleNamespace()) + assert tracking._wandb_usable() is True + + +def test_build_tracker_noops_cleanly_without_wandb(monkeypatch, tmp_path) -> None: + """A wandb-project config must degrade to the dummy tracker, not crash.""" + monkeypatch.setattr(tracking, "get_world_group", lambda: SimpleNamespace(rank=0)) + monkeypatch.setattr(tracking, "_wandb_usable", lambda: False) + + tracker = tracking.build_tracker( + TrackerConfig(project_name="h3-dmd2-vsa"), + CheckpointConfig(output_dir=str(tmp_path)), + config={"method": {}}, + ) + + assert isinstance(tracker, DummyTracker) + tracker.log({"total_loss": 1.0}, 1) # must not raise + tracker.finish() + + +def test_build_tracker_logs_loss_dict_with_monkeypatched_wandb(monkeypatch, tmp_path) -> None: + fake = _fake_wandb() + monkeypatch.setitem(sys.modules, "wandb", fake) + monkeypatch.setenv("WANDB_API_KEY", "key") + monkeypatch.setattr(tracking, "get_world_group", lambda: SimpleNamespace(rank=0)) + + run_config = {"method": {"dmd_denoising_steps": [1000, 757, 522]}} + tracker = tracking.build_tracker( + TrackerConfig(project_name="h3-dmd2-vsa", run_name="dmd2_vsa0_overfit"), + CheckpointConfig(output_dir=str(tmp_path)), + config=run_config, + ) + + assert isinstance(tracker, WandbTracker) + (run, ) = fake.runs + assert run.init_kwargs["project"] == "h3-dmd2-vsa" + assert run.init_kwargs["name"] == "dmd2_vsa0_overfit" + assert run.init_kwargs["config"] == run_config + + # The per-step dict the Trainer logs on rank 0 (DMD2 loss map + metrics). + metrics = { + "total_loss": 0.5, + "generator_loss": 0.25, + "fake_score_loss": 0.25, + "update_student": 1.0, + "step_time_sec": 0.1, + "vsa_sparsity": 0.0, + } + tracker.log(metrics, 7) + assert run.logged == [(metrics, 7)] + + +def test_build_tracker_nonzero_rank_never_inits_wandb(monkeypatch, tmp_path) -> None: + fake = _fake_wandb() + monkeypatch.setitem(sys.modules, "wandb", fake) + monkeypatch.setenv("WANDB_API_KEY", "key") + monkeypatch.setattr(tracking, "get_world_group", lambda: SimpleNamespace(rank=1)) + + tracker = tracking.build_tracker( + TrackerConfig(project_name="h3-dmd2-vsa"), + CheckpointConfig(output_dir=str(tmp_path)), + config=None, + ) + + assert isinstance(tracker, DummyTracker) + assert fake.runs == [] diff --git a/fastvideo/train/models/minimax_h3/minimax_h3.py b/fastvideo/train/models/minimax_h3/minimax_h3.py index b2eb86d62..018f5d44f 100644 --- a/fastvideo/train/models/minimax_h3/minimax_h3.py +++ b/fastvideo/train/models/minimax_h3/minimax_h3.py @@ -39,6 +39,15 @@ _AUDIO_SCHEDULER_SHIFT = 3.0 _VIDEO_LATENT_CHANNELS = 24 _AUDIO_LATENT_CHANNELS = 32 +# Dense TORCH_SDPA is the default; per-role overrides allow FLASH_ATTN +# (teacher/critic, FA4 via FASTVIDEO_FA4=1) and the packed-sequence VSA-H3 +# backend (distillation student). +_ALLOWED_ATTENTION_BACKENDS = ( + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3, +) + def shift_noise_amount(base_noise_amount: torch.Tensor, shift: float) -> torch.Tensor: """Apply the MiniMax H3 rational shift to a unit noise amount.""" @@ -68,10 +77,12 @@ class MiniMaxH3Model(ModelBase): trainable=trainable, attention_backend=attention_backend, ) - # PyTorch scaled dot product attention (SDPA) provides dense attention - # without adding another attention-kernel dependency to H3 training. - if self.attention_backend != AttentionBackendEnum.TORCH_SDPA: - raise ValueError("MiniMaxH3Model requires the TORCH_SDPA attention backend") + # Attention layers bind their backend during construction, so this is + # the per-role selection point (student/teacher/critic can differ). + if self.attention_backend not in _ALLOWED_ATTENTION_BACKENDS: + allowed = ", ".join(b.name for b in _ALLOWED_ATTENTION_BACKENDS) + raise ValueError("MiniMaxH3Model supports the attention backends " + f"{{{allowed}}}, got {self.attention_backend}") if training_config.pipeline_config is None: raise ValueError("MiniMaxH3Model requires a resolved MiniMax H3 pipeline config") # Packed row indices describe one text-video-audio document without a @@ -320,12 +331,12 @@ class MiniMaxH3Model(ModelBase): ) -> NoisePrediction: """Pack modality timesteps and convert H3 outputs to noise-minus-clean.""" del timestep - # Both attention-metadata views are None under TORCH_SDPA, so "vsa" - # silently means dense here (mirrors WanModel under dense backends). - # Real VSA-H3 metadata would slot in at this branch once the VSA-H3 - # attention backend lands (PR #1695). + # Under dense backends both metadata views are None, so "vsa" silently + # means dense (mirrors WanModel). MiniMaxH3DMDModel.prepare_batch + # populates attn_metadata_vsa when the role runs VIDEO_SPARSE_ATTN_H3. if attn_kind not in ("dense", "vsa"): raise ValueError(f"Unknown attn_kind: {attn_kind!r}") + attn_metadata = (batch.attn_metadata_vsa if attn_kind == "vsa" else batch.attn_metadata) layout = batch.minimax_h3_layout if not isinstance(layout, MiniMaxH3PackedLayout): raise RuntimeError("prepare_batch() must set TrainingBatch.minimax_h3_layout") @@ -366,7 +377,7 @@ class MiniMaxH3Model(ModelBase): with torch.autocast(device.type, dtype=dtype), set_forward_context( current_timestep=unique_timesteps, - attn_metadata=None, + attn_metadata=attn_metadata, ): video_velocity, audio_velocity = self.transformer( hidden_states=video_rows[None], diff --git a/fastvideo/train/models/minimax_h3/minimax_h3_dmd.py b/fastvideo/train/models/minimax_h3/minimax_h3_dmd.py index a58d7a217..2964ccf7d 100644 --- a/fastvideo/train/models/minimax_h3/minimax_h3_dmd.py +++ b/fastvideo/train/models/minimax_h3/minimax_h3_dmd.py @@ -8,11 +8,16 @@ from typing import Any, Literal import torch +import fastvideo.envs as envs +from fastvideo.attention.backends.video_sparse_attn_h3 import ( + MiniMaxH3VSAMetadataBuilder, ) from fastvideo.pipelines import TrainingBatch from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_AUDIO_CHANNELS, audio_latent_num_frames, ) +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import ( + _h3_vsa_prefix_segments, ) from fastvideo.train.models.minimax_h3.minimax_h3 import ( _AUDIO_LATENT_CHANNELS, _AUDIO_SCHEDULER_SHIFT, @@ -127,8 +132,40 @@ class MiniMaxH3DMDModel(MiniMaxH3Model): # clean latents matter here. The fine-tuning noisy fields are # refreshed by predict_noise on every call. batch.latents = self.pack_latents(batch.latents, batch.audio_latents) + self._maybe_build_vsa_metadata(batch) return batch + def _maybe_build_vsa_metadata(self, batch: TrainingBatch) -> None: + """Populate the VSA view when this role runs the VSA-H3 backend. + + DMD2 hardcodes student forwards to ``attn_kind="vsa"``; the dense + teacher/critic views stay ``None`` so their forwards remain dense. + One builder per model keeps the padded tile buffer reused across + steps (see ``MiniMaxH3VSAMetadata.tile_buf_holder``). + """ + backend = (self.attention_backend_name or envs.FASTVIDEO_ATTENTION_BACKEND) + if backend != "VIDEO_SPARSE_ATTN_H3": + return + layout = batch.minimax_h3_layout + patch_size = tuple(self.transformer.patch_size) + builder = getattr(self, "_vsa_metadata_builder", None) + if builder is None: + builder = self._vsa_metadata_builder = MiniMaxH3VSAMetadataBuilder() + batch.attn_metadata_vsa = builder.build( + # Training builds one metadata per batch and reuses it across the + # rollout's timesteps; the step index only feeds probe bookkeeping. + current_timestep=0, + raw_latent_shape=( + layout.num_video_latent_frames, + layout.latent_height, + layout.latent_width, + ), + patch_size=patch_size, + VSA_sparsity=float(self.training_config.vsa_sparsity), + prefix_segments=_h3_vsa_prefix_segments(layout, patch_size), + device=self.device, + ) + def add_noise( self, clean_latents: torch.Tensor, diff --git a/fastvideo/train/utils/tracking.py b/fastvideo/train/utils/tracking.py index 91e06889b..c552758c0 100644 --- a/fastvideo/train/utils/tracking.py +++ b/fastvideo/train/utils/tracking.py @@ -6,6 +6,7 @@ import os from typing import Any, TYPE_CHECKING from fastvideo.distributed import get_world_group +from fastvideo.logger import init_logger from fastvideo.training.trackers import ( initialize_trackers, Trackers, @@ -17,6 +18,25 @@ if TYPE_CHECKING: TrackerConfig, ) +logger = init_logger(__name__) + + +def _wandb_usable() -> bool: + """True when wandb can actually start a run (importable + credentials).""" + try: + import wandb + except Exception: + return False + if os.environ.get("WANDB_API_KEY"): + return True + if os.environ.get("WANDB_MODE") in ("offline", "disabled"): + return True + try: + # Covers ~/.netrc logins from a prior `wandb login`. + return wandb.api.api_key is not None + except Exception: + return False + def build_tracker( tracker_config: TrackerConfig, @@ -33,6 +53,11 @@ def build_tracker( trackers.append(Trackers.WANDB.value) if world_group.rank != 0: trackers = [] + if Trackers.WANDB.value in trackers and not _wandb_usable(): + logger.warning("wandb tracking requested but wandb is not importable " + "or no credentials are configured (WANDB_API_KEY); " + "continuing without the wandb tracker.") + trackers = [t for t in trackers if t != Trackers.WANDB.value] tracker_log_dir = (checkpoint_config.output_dir or os.getcwd()) if trackers: