[train] VSA-H3 DMD2 overfit: per-role backends, VSA metadata wiring, wandb guard
- MiniMaxH3Model: allow per-role attention backends (TORCH_SDPA/FLASH_ATTN/ VIDEO_SPARSE_ATTN_H3); predict_noise routes attn_metadata_vsa for attn_kind='vsa' (dense roles keep None). - MiniMaxH3DMDModel.prepare_batch builds real VSA-H3 metadata (packed text|cond|audio prefix segments via _h3_vsa_prefix_segments, sparsity from training.vsa.sparsity) when the role runs VIDEO_SPARSE_ATTN_H3. - VSA-H3 tile(): fresh buffer for grad-tracking forwards; shared-buffer reuse trips autograd's in-place version check at student backward. - build_tracker: degrade to DummyTracker when wandb is unimportable or has no credentials (modular stack already logs per-step loss dicts via the wandb tracker; no new callback needed). - dmd2_vsa0_overfit.yaml: student VSA-H3 @ sparsity 0.0, teacher/critic FLASH_ATTN, 3-step DMD schedule, wandb project h3-dmd2-vsa. - preprocess_minimax_h3_overfit: --video/--prompt/--output-dir/--model-path CLI; overfit_vsa0.sh: preprocess + 4-GPU torchrun with FASTVIDEO_FA4=1. - CPU tests: VSA metadata build/routing, per-role backend override, wandb guard + loss-dict logging, tile-under-grad regression.
This commit is contained in:
SolitaryThinker
@@ -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"
|
||||
@@ -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: {}
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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 == []
|
||||
@@ -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],
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user