diff --git a/examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_nvfp4_qad.yaml b/examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_nvfp4_qad.yaml new file mode 100644 index 000000000..582ccfd85 --- /dev/null +++ b/examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_nvfp4_qad.yaml @@ -0,0 +1,110 @@ +# Pika H3 Flex FL2VA NVFP4 QAD (the FastH3 QAD recipe on Flex's six-forward ladder). +# +# Teacher: Pika H3 Flex FL2VA in bf16, i.e. MiniMax-H3 `transformer` with the +# released `fl2va` adapter merged by +# scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py. Student: the same +# weights with NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks, +# static FFN scales from the Flex FL2VA calibration table, BF16 dense attention. +# +# Rows are text-only (T2VA) prompt manifests in the V2 QAD format +# (scripts/distill/minimax_h3_v2_qad/{build_v2_prompts,precompute_v2_text}.py); +# rows may carry their own width/height/num_frames on Flex's bucket grid. +# +# Paths are placeholders: the merged model, the calibration table (pin its +# sha256 once calibrated), the prompt manifests and the output dir. +models: + student: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexFL2VAModel + init_from: data/models/pika-flex-fl2va-merged + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: examples/train/calibration/pika_h3_flex/amax-fl2va.json + amax_sha256: null + # Attention linears use the unit activation scale; only the FFN + # linears read their calibrated amax from this table. + attention: true + ffn: true + skip_blocks: [] + fp4_attention: + enabled: false + data: + train_manifest: data/pika-flex-fl2va/train.jsonl + eval_manifest: data/pika-flex-fl2va/eval.jsonl + width: 832 + height: 480 + num_frames: 124 + teacher: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexFL2VAModel + init_from: data/models/pika-flex-fl2va-merged + trainable: false + frozen_bf16: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod + video_weight: 1.0 + audio_weight: 1.0 + rollout_fraction: 0.25 + rollout_start_step: 20 + # Half of the six forwards train per row, as OmniRef trains 4 of 8. + rungs_per_row: 3 + eval_every: 15 + full_eval_every: 30 + patience: 2 + fidelity_tolerance: 0.05 + require_bitwise_init: true + save_best: true + select_metric: score + eval_save_dir: outputs/pika_flex_fl2va_nvfp4_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 10 + noise_floor_eval: true + step0_cache: outputs/pika_flex_fl2va_nvfp4_qad/step0_cache.json + +training: + distributed: + num_gpus: 16 + sp_size: 4 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 16 + pin_cpu_memory: false + data: + train_batch_size: 1 + seed: 1234 + optimizer: + learning_rate: 5.0e-06 + betas: [0.9, 0.999] + eps: 1.0e-08 + weight_decay: 0.0 + lr_scheduler: cosine_with_min_lr + lr_warmup_steps: 10 + min_lr_ratio: 0.3 + loop: + max_train_steps: 200 + gradient_accumulation_steps: 1 + checkpoint: + output_dir: outputs/pika_flex_fl2va_nvfp4_qad + training_state_checkpointing_steps: 25 + checkpoints_total_limit: 2 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: pika-h3-flex-nvfp4 + run_name: pika-flex-fl2va-nvfp4-qad + model: + enable_gradient_checkpointing_type: full + dit_precision: fp32 + vsa_sparsity: 0.0 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + +pipeline: {} diff --git a/examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_nvfp4_qad.yaml b/examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_nvfp4_qad.yaml new file mode 100644 index 000000000..3cebc4784 --- /dev/null +++ b/examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_nvfp4_qad.yaml @@ -0,0 +1,131 @@ +# Pika H3 Flex Ref2VA NVFP4 QAD (the FastH3 QAD recipe on Flex's six-forward ladder). +# +# Teacher: Pika H3 Flex Ref2VA in bf16, i.e. MiniMax-H3 `transformer_ref` with the +# released `ref2va` adapter merged by +# scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py. Student: the same +# weights with NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks +# active from step 0, static FFN activation scales from the Flex calibration +# table. Attention stays BF16 dense: Flex has no VSA gate, so there is no FP4 +# sparse-attention stage and no gate to freeze. Norms, AdaLN, the token refiner +# and the FP32 boundary modules stay frozen. +# +# The objective, evaluation, guards and checkpoint contract are the shared +# recipe's (docs/training/minimax_h3_fp4_qad.md). Rows are OmniRef Ref2VA +# parquet rows; the teacher computes every target on the fly, so the OmniRef +# held-out eval manifest is reused as is. +# +# Paths are placeholders: the merged model, the calibration table (pin its +# sha256 once calibrated), the row manifests and the output dir. +models: + student: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel + init_from: data/models/pika-flex-ref2va-merged + trainable: true + enable_gradient_checkpointing_type: full + qad: + enabled: true + linears: + enabled: true + amax_json: examples/train/calibration/pika_h3_flex/amax-ref2va.json + amax_sha256: null + # Attention linears use the unit activation scale; only the FFN + # linears read their calibrated amax from this table. + attention: true + ffn: true + skip_blocks: [] + fp4_attention: + enabled: false + data: + manifests: + - data/omniref/train-manifest.jsonl + cases: + - first_frame + - first_last_frame + - storyboard + - continue_scene + - continue_shot + - full_scene_audio_image_reference + per_group: + 480p: 400 + 768p: 50 + seed: 20261007 + max_frames: 243 + eval_manifest: data/omniref/eval-manifest.json + eval_resolutions: + - 480p + - 768p + eval_per_group: + 480p: 2 + 768p: 1 + eval_max_frames: 0 + teacher: + _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel + init_from: data/models/pika-flex-ref2va-merged + trainable: false + frozen_bf16: true + +method: + _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod + video_weight: 1.0 + audio_weight: 1.0 + rollout_fraction: 0.25 + rollout_start_step: 20 + # Half of the six forwards train per row, as OmniRef trains 4 of 8. + rungs_per_row: 3 + eval_every: 15 + full_eval_every: 30 + patience: 2 + fidelity_tolerance: 0.05 + require_bitwise_init: true + save_best: true + select_metric: score + eval_save_dir: outputs/pika_flex_ref2va_nvfp4_qad/samples + objective: output+recon + recon_lambda: 1.0 + recon_normalize: step0 + max_rollout_rows: 10 + noise_floor_eval: true + step0_cache: outputs/pika_flex_ref2va_nvfp4_qad/step0_cache.json + +training: + distributed: + num_gpus: 16 + sp_size: 4 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 16 + pin_cpu_memory: false + data: + train_batch_size: 1 + seed: 1234 + optimizer: + learning_rate: 5.0e-06 + betas: [0.9, 0.999] + eps: 1.0e-08 + weight_decay: 0.0 + lr_scheduler: cosine_with_min_lr + lr_warmup_steps: 10 + min_lr_ratio: 0.3 + loop: + max_train_steps: 200 + gradient_accumulation_steps: 1 + checkpoint: + output_dir: outputs/pika_flex_ref2va_nvfp4_qad + training_state_checkpointing_steps: 25 + checkpoints_total_limit: 2 + resume_from_checkpoint: '' + tracker: + trackers: [wandb] + project_name: pika-h3-flex-nvfp4 + run_name: pika-flex-ref2va-nvfp4-qad + model: + enable_gradient_checkpointing_type: full + dit_precision: fp32 + vsa_sparsity: 0.0 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + +pipeline: {} diff --git a/fastvideo/tests/train/models/test_pika_h3_flex_qad.py b/fastvideo/tests/train/models/test_pika_h3_flex_qad.py new file mode 100644 index 000000000..cd5201460 --- /dev/null +++ b/fastvideo/tests/train/models/test_pika_h3_flex_qad.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks for the Pika H3 Flex QAD profile: ladder, re-noise determinism, contract, adapter merge.""" +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import scheduler_step +from fastvideo.train.models.minimax_h3.flex import (CONTRACT_FILE, CONTRACT_SCHEMA, FlexContract, + MiniMaxH3FlexRef2VAModel, RowRenoiseScheduler, flex_sigmas, + renoise_seed) + +# Pika's published per-modality ladders (knots 1000, 750, 500, 250, 100, 40, 0 shifted by 12 / 3). +PIKA_VIDEO_SIGMAS = (1.0, 0.97297, 0.92308, 0.8, 0.57143, 0.33333, 0.0) +PIKA_AUDIO_SIGMAS = (1.0, 0.9, 0.75, 0.5, 0.25, 0.11111, 0.0) +STEPS = (1000, 750, 500, 250, 100, 40) + + +def _merge_module(): + path = Path(__file__).resolve().parents[4] / "scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py" + spec = importlib.util.spec_from_file_location("merge_pika_adapter", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _CPUFlexModel(MiniMaxH3FlexRef2VAModel): + device = torch.device("cpu") + + +def _model_stub() -> MiniMaxH3FlexRef2VAModel: + model = object.__new__(_CPUFlexModel) + model.contract = FlexContract("ref2va", "transformer_ref", STEPS, 12.0, 3.0) + return model + + +def _row(seed: int) -> SimpleNamespace: + return SimpleNamespace(spec=SimpleNamespace(seed=seed)) + + +def test_flex_sigmas_match_pika_ladders(): + torch.testing.assert_close(flex_sigmas(STEPS, 12.0), torch.tensor(PIKA_VIDEO_SIGMAS), atol=1e-5, rtol=0) + torch.testing.assert_close(flex_sigmas(STEPS, 3.0), torch.tensor(PIKA_AUDIO_SIGMAS), atol=1e-5, rtol=0) + + +def test_schedulers_expose_six_forwards_in_clean_time(): + video, audio = _model_stub().schedulers(_row(7)) + assert video.timesteps.numel() == audio.timesteps.numel() == 6 + torch.testing.assert_close(video.timesteps, 1 - torch.tensor(PIKA_VIDEO_SIGMAS[:-1]), atol=1e-5, rtol=0) + assert video.stochastic_sampling and audio.stochastic_sampling + + +def test_schedulers_require_the_row(): + with pytest.raises(ValueError, match="prepared row"): + _model_stub().schedulers() + + +def test_renoise_step_is_x0_plus_fixed_fresh_noise(): + video, _ = _model_stub().schedulers(_row(11)) + state, velocity = torch.randn(5, 4), torch.randn(5, 4) + rung = 2 + out = scheduler_step(video, rung, velocity, state) + timestep = video.timesteps[rung] + x0 = state + (1 - timestep) * velocity + generator = torch.Generator().manual_seed(renoise_seed(11, rung, 0)) + noise = torch.randn(state.shape, generator=generator, dtype=torch.float32) + sigma_next = video.sigmas[rung + 1] + torch.testing.assert_close(out, (1 - sigma_next) * x0 + sigma_next * noise) + + +def test_renoise_noise_depends_only_on_row_forward_and_modality(): + """Teacher-forced pass, rollout and re-evaluation of one row must draw identical noise at each forward.""" + state, velocity = torch.randn(6, 3), torch.randn(6, 3) + first = _model_stub().schedulers(_row(3)) + second = _model_stub().schedulers(_row(3)) + # Out-of-order use (a rollout after a full teacher-forced pass) still reproduces forward 1. + for rung in range(5): + scheduler_step(first[0], rung, velocity, state) + torch.testing.assert_close(scheduler_step(first[0], 1, velocity, state), + scheduler_step(second[0], 1, velocity, state)) + assert not torch.equal(scheduler_step(first[0], 1, velocity, state), scheduler_step(first[0], 2, velocity, + state)) + other_row = _model_stub().schedulers(_row(4)) + assert not torch.equal(scheduler_step(first[0], 1, velocity, state), + scheduler_step(other_row[0], 1, velocity, state)) + assert renoise_seed(3, 1, 0) != renoise_seed(3, 1, 1) + + +def test_last_forward_returns_x0_without_noise(): + video, audio = _model_stub().schedulers(_row(5)) + state, velocity = torch.randn(4, 2), torch.randn(4, 2) + for scheduler in (video, audio): + out = scheduler_step(scheduler, 5, velocity, state) + torch.testing.assert_close(out, state + (1 - scheduler.timesteps[5]) * velocity) + + +def test_row_renoise_scheduler_refuses_an_implicit_index(): + scheduler = RowRenoiseScheduler(12.0, 0, 0) + scheduler.set_timesteps(sigmas=flex_sigmas(STEPS, 12.0)) + with pytest.raises(ValueError, match="explicit step index"): + scheduler.step(torch.zeros(2, 2), scheduler.timesteps[0], torch.zeros(2, 2)) + + +def _write_contract(path: Path, **overrides) -> None: + data = { + "schema": CONTRACT_SCHEMA, + "task": "ref2va", + "trunk": "transformer_ref", + "denoising_steps": list(STEPS), + "video_scheduler_shift": 12.0, + "audio_scheduler_shift": 3.0, + "sampler": "renoise", + } + data.update(overrides) + (path / CONTRACT_FILE).write_text(json.dumps(data)) + + +def test_contract_reads_the_released_ladder(tmp_path): + _write_contract(tmp_path) + assert FlexContract.read(str(tmp_path)) == FlexContract("ref2va", "transformer_ref", STEPS, 12.0, 3.0) + + +@pytest.mark.parametrize("overrides", [{ + "denoising_steps": [1000, 500, 0] +}, { + "video_scheduler_shift": 10.0 +}, { + "trunk": "transformer" +}, { + "sampler": "euler" +}]) +def test_contract_rejects_anything_but_the_release(tmp_path, overrides): + _write_contract(tmp_path, **overrides) + with pytest.raises(ValueError): + FlexContract.read(str(tmp_path)) + + +def test_contract_must_match_the_profile_trunk(tmp_path): + _write_contract(tmp_path, task="fl2va", trunk="transformer") + with pytest.raises(ValueError, match="needs ref2va on transformer_ref"): + _model_stub()._read_contract(str(tmp_path)) + + +def test_fp4_attention_is_refused(): + model = _model_stub() + model.qad_config = {"fp4_attention": {"enabled": True}, "linears": {"enabled": True}} + with pytest.raises(ValueError, match="dense attention"): + model._install_qad() + + +def _write_trunk(trunk: Path, weights: dict[str, torch.Tensor]) -> None: + trunk.mkdir(parents=True) + names = sorted(weights) + shards = {"model-00001-of-00002.safetensors": names[:2], "model-00002-of-00002.safetensors": names[2:]} + weight_map = {} + for shard, keys in shards.items(): + save_file({key: weights[key] for key in keys}, str(trunk / shard)) + weight_map.update({key: shard for key in keys}) + (trunk / "diffusion_pytorch_model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + (trunk / "config.json").write_text("{}") + + +def test_merge_trunk_adds_scaled_lora_in_fp32(tmp_path): + merge = _merge_module() + torch.manual_seed(0) + weights = { + "blocks.0.attn.to_q.weight": torch.randn(6, 4).to(torch.bfloat16), + "blocks.0.attn.to_q.bias": torch.randn(6).to(torch.bfloat16), + "blocks.0.ff.net.2.weight": torch.randn(4, 8).to(torch.bfloat16), + "proj_out.weight": torch.randn(3, 4), + } + factors = { + "blocks.0.attn.to_q": (torch.randn(2, 4), torch.randn(6, 2)), + "blocks.0.ff.net.2": (torch.randn(2, 8), torch.randn(4, 2)), + } + _write_trunk(tmp_path / "base", weights) + report = merge.merge_trunk(tmp_path / "base", tmp_path / "out", factors, scale=0.5) + assert report["modules"] == 2 + from safetensors.torch import load_file + merged = {} + for shard in (tmp_path / "out").glob("*.safetensors"): + merged.update(load_file(str(shard))) + for module, (lora_a, lora_b) in factors.items(): + key = f"{module}.weight" + expected = (weights[key].float() + 0.5 * lora_b @ lora_a).to(torch.bfloat16) + assert merged[key].dtype == torch.bfloat16 + assert torch.equal(merged[key], expected) + assert torch.equal(merged["blocks.0.attn.to_q.bias"], weights["blocks.0.attn.to_q.bias"]) + assert torch.equal(merged["proj_out.weight"], weights["proj_out.weight"]) + + +def test_merge_trunk_refuses_adapter_modules_without_weights(tmp_path): + merge = _merge_module() + _write_trunk(tmp_path / "base", {f"w{i}.weight": torch.zeros(2, 2) for i in range(3)}) + with pytest.raises(ValueError, match="no weight"): + merge.merge_trunk(tmp_path / "base", tmp_path / "out", {"missing": (torch.zeros(1, 2), torch.zeros(2, 1))}, + scale=1.0) + + +def test_adapter_factors_pair_peft_keys(tmp_path): + merge = _merge_module() + path = tmp_path / "adapter.safetensors" + save_file( + { + "a.to_q.lora_A.default.weight": torch.zeros(2, 4), + "a.to_q.lora_B.default.weight": torch.zeros(6, 2), + }, str(path)) + factors = merge.adapter_factors(str(path)) + assert list(factors) == ["a.to_q"] and factors["a.to_q"][0].shape == (2, 4) + save_file({"a.to_q.lora_A.default.weight": torch.zeros(2, 4)}, str(path)) + with pytest.raises(ValueError, match="both LoRA factors"): + merge.adapter_factors(str(path)) diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad.py index d605792a2..5e5588136 100644 --- a/fastvideo/train/methods/knowledge_distillation/pdd_qad.py +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad.py @@ -275,7 +275,7 @@ class OmniRefPDDQADMethod(TrainingMethod): spec = batch["row"] prepared = self._row(spec) rollout = self._use_rollout(spec.id, iteration) - schedulers = self.student.schedulers() + schedulers = self.student.schedulers(prepared) timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) accum = max(1, int(self.training_config.loop.gradient_accumulation_steps or 1)) num_rungs = self.student.num_rungs @@ -355,7 +355,7 @@ class OmniRefPDDQADMethod(TrainingMethod): per_rank = -(-self.recon_norm_rows // self.dp_size) for spec, counted in rows[:per_rank]: prepared = self._row(spec) - schedulers = self.student.schedulers() + schedulers = self.student.schedulers(prepared) timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) video, audio = prepared.video, prepared.audio for rung in range(num_rungs): @@ -571,7 +571,7 @@ class OmniRefPDDQADMethod(TrainingMethod): for position, (spec, counted) in enumerate(rows): stats = counted_stats if counted else RungStats(num_rungs) prepared = self._row(spec) - schedulers = self.student.schedulers() + schedulers = self.student.schedulers(prepared) timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) video, audio = prepared.video, prepared.audio for rung in range(num_rungs): @@ -651,7 +651,7 @@ class OmniRefPDDQADMethod(TrainingMethod): diff = torch.zeros((), device=self.student.device, dtype=torch.float64) if rows: prepared = self._row(rows[0][0]) - schedulers = self.student.schedulers() + schedulers = self.student.schedulers(prepared) timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) with torch.no_grad(), self.student.quantization(linears=False, attention=False): for rung in (0, self.student.num_rungs - 1): @@ -699,7 +699,7 @@ class OmniRefPDDQADMethod(TrainingMethod): with self.student.quantization(linears=False, attention=False), calibrating(self.student.qad_linears): for spec in mine: prepared = self._row(spec) - schedulers = self.student.schedulers() + schedulers = self.student.schedulers(prepared) timesteps = (schedulers[0].timesteps, schedulers[1].timesteps) video, audio = prepared.video, prepared.audio for rung in range(self.student.num_rungs): diff --git a/fastvideo/train/models/minimax_h3/flex.py b/fastvideo/train/models/minimax_h3/flex.py new file mode 100644 index 000000000..5a8fc35e7 --- /dev/null +++ b/fastvideo/train/models/minimax_h3/flex.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pika H3 Flex (six-forward re-noise sampler) as a per-forward model for NVFP4 QAD. + +Same contract as ``omniref_pdd.MiniMaxH3OmniRefPDDModel``, so ``OmniRefPDDQADMethod`` runs unchanged: one rung per +sampler forward, ``rung_scope`` / ``forward_rung`` / ``schedulers`` / ``prepare`` / ``quantization``. ``init_from`` +is a merged Flex model directory written by ``scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py``. + +Differences from OmniRef: + +- Dense attention (``FLASH_ATTN``): Flex has no VSA compression gate, so there is no tile selection to freeze and the + FP4 sparse-attention emulator does not apply. QAD covers the NVFP4 linears only. +- No PDD head fusion: every rung is a plain transformer forward over Pika's clean-time knots + (1000, 750, 500, 250, 100, 40), shifted per modality (video 12, audio 3), starting from unit noise. +- Re-noise steps: ``x_next = (1 - sigma_next) x0 + sigma_next eps``. ``eps`` is a function of (row seed, forward, + modality) only, so the teacher-forced pass, the student rollout and the seed+1 noise-floor eval of a row all draw + the same noise. It is not Pika's sampler stream; the distribution is. + +``MiniMaxH3FlexRef2VAModel`` reads OmniRef Ref2VA rows (``omniref_data``) on ``transformer_ref``; +``MiniMaxH3FlexFL2VAModel`` reads text-only rows (``v2_data``) on ``transformer``. +""" +from __future__ import annotations + +import contextlib +import json +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch + +from fastvideo.configs.pipelines.pika_h3_flex import (PIKA_AUDIO_SCHEDULER_SHIFT, PIKA_DENOISING_STEPS, + PIKA_VIDEO_SCHEDULER_SHIFT) +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler +from fastvideo.train.models.minimax_h3.omniref_pdd import MiniMaxH3OmniRefPDDModel, _pin_amax +from fastvideo.train.models.minimax_h3.v2_t2v import MiniMaxH3V2T2VModel + +if TYPE_CHECKING: + from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow + from fastvideo.train.utils.training_config import TrainingConfig + +logger = init_logger(__name__) +CONTRACT_FILE = "pika_flex_qad.json" +CONTRACT_SCHEMA = "pika-flex-qad-v1" +TASK_TO_TRUNK = {"fl2va": "transformer", "ref2va": "transformer_ref"} +# Block projections and FFN linears: the set the NVFP4 export quantizes. Norms, AdaLN, the token refiner, +# embeddings and the output heads stay frozen. +FLEX_TRAINABLE = r"transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out)|ff\.(fc_in|fc_out))\." +_VIDEO_STREAM, _AUDIO_STREAM = 0, 1 + + +@dataclass(frozen=True) +class FlexContract: + """The merged checkpoint's sampling contract (``pika_flex_qad.json``).""" + task: str + trunk: str + denoising_steps: tuple[int, ...] + video_shift: float + audio_shift: float + + @classmethod + def read(cls, model_path: str) -> FlexContract: + path = Path(model_path) / CONTRACT_FILE + if not path.is_file(): + raise ValueError(f"{model_path} has no {CONTRACT_FILE}; build it with merge_pika_adapter.py") + data = json.loads(path.read_text()) + if data.get("schema") != CONTRACT_SCHEMA or data.get("sampler") != "renoise": + raise ValueError(f"{path} is not a {CONTRACT_SCHEMA} re-noise contract") + contract = cls(str(data["task"]), str(data["trunk"]), tuple(int(s) for s in data["denoising_steps"]), + float(data["video_scheduler_shift"]), float(data["audio_scheduler_shift"])) + released = (PIKA_DENOISING_STEPS, PIKA_VIDEO_SCHEDULER_SHIFT, PIKA_AUDIO_SCHEDULER_SHIFT) + if (contract.denoising_steps, contract.video_shift, contract.audio_shift) != released: + raise ValueError(f"{path} does not carry the released Pika H3 Flex ladder {released}") + if TASK_TO_TRUNK.get(contract.task) != contract.trunk: + raise ValueError(f"{path}: task {contract.task!r} does not run on trunk {contract.trunk!r}") + return contract + + +def flex_sigmas(steps: tuple[int, ...], shift: float) -> torch.Tensor: + """Shifted sigma grid over ``steps`` plus the terminal zero, as ``MiniMaxH3DenoisingStage._set_shifted_schedule``.""" + base = torch.tensor([step / 1000.0 for step in steps] + [0.0], dtype=torch.float32) + return shift * base / (1 + (shift - 1) * base) + + +def renoise_seed(row_seed: int, forward: int, stream: int) -> int: + """Distinct generator seed per (row seed, sampler forward, modality); forwards < 8, two streams.""" + return (int(row_seed) + 1) * 16 + 2 * int(forward) + int(stream) + + +class RowRenoiseScheduler(MiniMaxH3Scheduler): + """H3 re-noise scheduler whose fresh noise at forward ``k`` is fixed by (row seed, ``k``, modality).""" + + def __init__(self, shift: float, row_seed: int, stream: int) -> None: + super().__init__(shift=shift, stochastic_sampling=True, shift_provided_sigmas=True) + self._row_seed = int(row_seed) + self._stream = int(stream) + + def step(self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor, + sample: torch.FloatTensor, + return_dict: bool = True) -> Any: + if self._step_index is None: + raise ValueError("RowRenoiseScheduler needs an explicit step index (use pdd_qad_metrics.scheduler_step)") + generator = torch.Generator(device=sample.device) + generator.manual_seed(renoise_seed(self._row_seed, self._step_index, self._stream)) + self.set_noise_generator(generator) + return super().step(model_output, timestep, sample, return_dict=return_dict) + + +class MiniMaxH3FlexRef2VAModel(MiniMaxH3OmniRefPDDModel): + """Teacher (merged bf16 Flex) or NVFP4 QAD student on the six-forward Flex ladder, Ref2VA rows.""" + + _transformer_module_type = "transformer_ref" + _attention_backend = "FLASH_ATTN" + _task = "ref2va" + + def __init__(self, + *, + init_from: str, + training_config: TrainingConfig, + trainable: bool = True, + enable_gradient_checkpointing_type: str | None = None, + trainable_patterns: str = FLEX_TRAINABLE, + qad: dict[str, Any] | None = None, + data: dict[str, Any] | None = None, + frozen_bf16: bool = False) -> None: + super().__init__(init_from=init_from, + training_config=training_config, + trainable=trainable, + enable_gradient_checkpointing_type=enable_gradient_checkpointing_type, + trainable_patterns=trainable_patterns, + qad=qad, + data=data, + frozen_bf16=frozen_bf16) + + # ------------------------------------------------------------------ construction hooks + def _read_contract(self, model_path: str) -> Any: + contract = FlexContract.read(model_path) + if contract.task != self._task or contract.trunk != self._transformer_module_type: + raise ValueError(f"{model_path} is a Flex {contract.task} merge on {contract.trunk}; " + f"{type(self).__name__} needs {self._task} on {self._transformer_module_type}") + return contract + + def _check_tile(self) -> None: + """Dense attention: no VSA tile to check.""" + + def _install_qad(self) -> None: + """NVFP4 W4A4 linears with the calibrated static FFN scales; attention stays BF16 dense.""" + from fastvideo.layers.quantization.nvfp4_qad import NVFP4QADPlan, install_nvfp4_qad, load_amax_table + + if dict(self.qad_config.get("fp4_attention") or {}).get("enabled", False): + raise ValueError("Pika H3 Flex runs dense attention; the FP4 attention emulator models VSA-H3 tiles. " + "Set qad.fp4_attention.enabled: false") + linears = dict(self.qad_config.get("linears") or {}) + if not linears.get("enabled", True): + return + amax_path = linears.get("amax_json") + table = load_amax_table(amax_path) if amax_path else None + if amax_path: + self.amax_sha256 = _pin_amax(amax_path, linears.get("amax_sha256"), + self.training_config.checkpoint.output_dir) + self.provisional_rows = int(linears.get("provisional_calibration_rows", 0)) + if table is None and self.provisional_rows <= 0: + raise ValueError("qad.linears.amax_json (the Flex calibration table) is required") + # Flex checkpoints have no compression gate, so ``gate`` has nothing to quantize. + plan = NVFP4QADPlan(quantize_attention=bool(linears.get("attention", True)), + quantize_gate=False, + quantize_ffn=bool(linears.get("ffn", True)), + skip_blocks=tuple(int(b) for b in linears.get("skip_blocks", ()))) + self.qad_linears = install_nvfp4_qad(self.transformer, table, plan, allow_uncalibrated=table is None) + logger.info("Flex %s QAD student: %d NVFP4 QAD linears, BF16 dense attention, %.2fB trainable parameters", + self._task, len(self.qad_linears), self.num_trainable / 1e9) + + @contextlib.contextmanager + def quantization(self, *, linears: bool = True, attention: bool = True) -> Iterator[None]: + """Toggle the NVFP4 linears (``attention`` is accepted for the method's contract; it is always BF16).""" + from fastvideo.layers.quantization.nvfp4_qad import set_nvfp4_qad_enabled + + set_nvfp4_qad_enabled(self.qad_linears, linears) + try: + yield + finally: + set_nvfp4_qad_enabled(self.qad_linears, True) + + # ------------------------------------------------------------------ Flex ladder + @property + def plan(self) -> Any: # pragma: no cover - PDD-only + raise AttributeError("Pika H3 Flex has no PDD plan") + + @property + def num_rungs(self) -> int: + return len(self.contract.denoising_steps) + + def schedulers(self, prepared: PreparedRow | None = None) -> tuple[Any, Any]: + """Re-noise schedulers on the shifted Flex grids; their noise is fixed by ``prepared``'s seed.""" + if prepared is None: + raise ValueError("Flex re-noise schedulers need the prepared row (its seed fixes the noise)") + seed = int(prepared.spec.seed) + out = [] + for shift, stream in ((self.contract.video_shift, _VIDEO_STREAM), (self.contract.audio_shift, _AUDIO_STREAM)): + scheduler = RowRenoiseScheduler(shift, seed, stream) + scheduler.set_timesteps(sigmas=flex_sigmas(self.contract.denoising_steps, shift), device=self.device) + out.append(scheduler) + return out[0], out[1] + + def prepare(self, spec: Any, row: dict[str, Any]) -> PreparedRow: + from fastvideo.train.models.minimax_h3.omniref_data import prepare_row + + # Knot 1000 is sigma 1 for both modalities: target rows start as unit noise (no PDD start scaling). + return prepare_row(spec, + row, + patch_size=self.patch_size, + latent_channels=24, + audio_channels=32, + pdd_start_sigmas=(1.0, 1.0), + device=self.device) + + @contextlib.contextmanager + def rung_scope(self, prepared: PreparedRow, rung: int, timesteps: tuple[Any, Any]) -> Iterator[Any]: + """Forward context of Flex forward ``rung``; keep backward inside it (activation recompute reads it).""" + from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG, build_row_timesteps + video_t, audio_t = float(timesteps[0][rung].item()), float(timesteps[1][rung].item()) + unique, inverse = build_row_timesteps(prepared.layout, + video_timestep=video_t, + audio_timestep=audio_t, + condition_video_timestep=max(video_t, MINIMAX_H3_KEYFRAME_NOISE_AUG), + condition_audio_timestep=1.0) + with set_forward_context(current_timestep=rung, attn_metadata=None): + yield unique.to(self.device), inverse.to(self.device) + + +class MiniMaxH3FlexFL2VAModel(MiniMaxH3FlexRef2VAModel): + """The FL2VA adapter on ``transformer``: text-only (T2VA) rows from ``v2_data`` manifests.""" + + _transformer_module_type = "transformer" + _task = "fl2va" + + # The data path is V2's (``v2_data`` manifests: prompt, Qwen3-VL embedding, seed); only the ladder differs. + init_preprocessors = MiniMaxH3V2T2VModel.init_preprocessors + load_row = MiniMaxH3V2T2VModel.load_row + + +__all__ = [ + "FLEX_TRAINABLE", "FlexContract", "MiniMaxH3FlexFL2VAModel", "MiniMaxH3FlexRef2VAModel", "RowRenoiseScheduler", + "flex_sigmas", "renoise_seed" +] diff --git a/fastvideo/train/models/minimax_h3/omniref_pdd.py b/fastvideo/train/models/minimax_h3/omniref_pdd.py index ca5bb436b..4cf45e551 100644 --- a/fastvideo/train/models/minimax_h3/omniref_pdd.py +++ b/fastvideo/train/models/minimax_h3/omniref_pdd.py @@ -70,6 +70,7 @@ class MiniMaxH3OmniRefPDDModel(ModelBase): _transformer_cls_name = "MiniMaxH3Transformer3DModel" _transformer_module_type = "transformer_ref" + _attention_backend = "VIDEO_SPARSE_ATTN_H3" def __init__(self, *, @@ -81,7 +82,7 @@ class MiniMaxH3OmniRefPDDModel(ModelBase): qad: dict[str, Any] | None = None, data: dict[str, Any] | None = None, frozen_bf16: bool = False) -> None: - super().__init__(trainable=trainable, attention_backend="VIDEO_SPARSE_ATTN_H3") + super().__init__(trainable=trainable, attention_backend=self._attention_backend) # Keep a frozen role's masters in bf16 (the checkpoint's dtype) instead of an fp32 upcast: compute is bf16 either # way, so outputs are bit-identical, and the role needs half the memory. self._frozen_bf16 = bool(frozen_bf16) @@ -284,7 +285,7 @@ class MiniMaxH3OmniRefPDDModel(ModelBase): def num_rungs(self) -> int: return int(self.plan.num_steps) - def schedulers(self) -> tuple[Any, Any]: + def schedulers(self, prepared: PreparedRow | None = None) -> tuple[Any, Any]: """Fresh video/audio schedulers on the PDD node sigmas (as the denoising stage sets them).""" from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler video = MiniMaxH3Scheduler(shift=self.contract.video_shift) diff --git a/fastvideo/train/models/minimax_h3/v2_t2v.py b/fastvideo/train/models/minimax_h3/v2_t2v.py index 98a093fd6..c4550abe7 100644 --- a/fastvideo/train/models/minimax_h3/v2_t2v.py +++ b/fastvideo/train/models/minimax_h3/v2_t2v.py @@ -189,7 +189,7 @@ class MiniMaxH3V2T2VModel(MiniMaxH3OmniRefPDDModel): def num_rungs(self) -> int: return len(self.contract.dmd_steps) - def schedulers(self) -> tuple[Any, Any]: + def schedulers(self, prepared: PreparedRow | None = None) -> tuple[Any, Any]: """Fresh video/audio schedulers on the DMD sigma grid, as ``MiniMaxH3DenoisingStage._set_dmd_schedule``.""" from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler base = torch.tensor([step / 1000.0 for step in self.contract.dmd_steps] + [0.0], dtype=torch.float32) diff --git a/scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py b/scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py new file mode 100644 index 000000000..cdb192f14 --- /dev/null +++ b/scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Merge one released Pika H3 Flex adapter into its MiniMax-H3 trunk for QAD. + +QAD needs the Flex teacher as plain weights: the student's fp32 masters start +from them and the NVFP4 export quantizes them. This writes a model directory +whose trunk (``transformer`` for FL2VA, ``transformer_ref`` for Ref2VA) holds +``W + (alpha / rank) * B @ A`` for every adapted module, merged in fp32 and +rounded once to the base dtype, so the result does not depend on the runtime +LoRA path. Every other component is a symlink to the base snapshot. + +The adapter is located and validated with the inference loader +(``fastvideo.pipelines.basic.pika_h3_flex.adapter.load_adapter``); its sha256 +must match the manifest. The output also carries ``pika_flex_qad.json``, the +sampling contract the QAD model profile reads. + + python scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py \\ + --base --adapter Pika-Labs/Pika-H3-Flex \\ + --task ref2va --output data/models/pika-flex-ref2va-merged +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from fastvideo.configs.pipelines.pika_h3_flex import (PIKA_AUDIO_SCHEDULER_SHIFT, PIKA_DENOISING_STEPS, + PIKA_VIDEO_SCHEDULER_SHIFT) +from fastvideo.pipelines.basic.pika_h3_flex.adapter import load_adapter + +CONTRACT_FILE = "pika_flex_qad.json" +CONTRACT_SCHEMA = "pika-flex-qad-v1" +_LORA_A = ".lora_A.default.weight" +_LORA_B = ".lora_B.default.weight" + + +def adapter_factors(path: str) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: + """``{module: (A [r, in], B [out, r])}`` from a PEFT adapter file.""" + factors: dict[str, dict[str, torch.Tensor]] = {} + with safe_open(path, framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a mapping + for suffix, slot in ((_LORA_A, "A"), (_LORA_B, "B")): + if key.endswith(suffix): + factors.setdefault(key[:-len(suffix)], {})[slot] = handle.get_tensor(key) + break + else: + raise ValueError(f"{path}: unexpected adapter tensor {key!r}") + incomplete = sorted(module for module, pair in factors.items() if set(pair) != {"A", "B"}) + if incomplete: + raise ValueError(f"{path}: modules without both LoRA factors: {incomplete[:5]}") + return {module: (pair["A"], pair["B"]) for module, pair in factors.items()} + + +def merge_weight(weight: torch.Tensor, lora_a: torch.Tensor, lora_b: torch.Tensor, scale: float) -> torch.Tensor: + """``weight + scale * B @ A`` in fp32, rounded once to ``weight``'s dtype.""" + if lora_b.shape[1] != lora_a.shape[0] or (lora_b.shape[0], lora_a.shape[1]) != tuple(weight.shape): + raise ValueError(f"LoRA shapes A {tuple(lora_a.shape)} / B {tuple(lora_b.shape)} do not fit weight " + f"{tuple(weight.shape)}") + delta = lora_b.to(torch.float32) @ lora_a.to(torch.float32) + return (weight.to(torch.float32) + scale * delta).to(weight.dtype) + + +def merge_trunk(trunk_dir: Path, output_dir: Path, factors: dict[str, tuple[torch.Tensor, torch.Tensor]], + scale: float) -> dict[str, Any]: + """Write ``trunk_dir``'s shards with every adapted module merged; return a per-module report.""" + index_path = trunk_dir / "diffusion_pytorch_model.safetensors.index.json" + index = json.loads(index_path.read_text()) + shards = sorted(set(index["weight_map"].values())) + output_dir.mkdir(parents=True, exist_ok=True) + remaining = set(factors) + report: dict[str, float] = {} + for shard in shards: + tensors: dict[str, torch.Tensor] = {} + with safe_open(str(trunk_dir / shard), framework="pt") as handle: + metadata = handle.metadata() + for key in handle.keys(): # noqa: SIM118 + tensor = handle.get_tensor(key) + module = key[:-len(".weight")] if key.endswith(".weight") else None + if module in factors: + lora_a, lora_b = factors[module] + merged = merge_weight(tensor, lora_a, lora_b, scale) + report[module] = float((merged.float() - tensor.float()).norm() / tensor.float().norm()) + remaining.discard(module) + tensor = merged + tensors[key] = tensor.contiguous() + save_file(tensors, str(output_dir / shard), metadata=metadata) + if remaining: + raise ValueError(f"{len(remaining)} adapter modules have no weight in {trunk_dir}: " + f"{sorted(remaining)[:5]}") + (output_dir / index_path.name).write_text(index_path.read_text()) + (output_dir / "config.json").write_text((trunk_dir / "config.json").read_text()) + return {"modules": len(report), "max_rel_delta": max(report.values()), "rel_delta": report} + + +def link_base(base: Path, output: Path, trunk: str) -> None: + """Symlink every base entry except the merged trunk.""" + output.mkdir(parents=True, exist_ok=True) + for entry in sorted(base.iterdir()): + if entry.name == trunk: + continue + target = output / entry.name + if target.is_symlink() or target.exists(): + raise FileExistsError(f"{target} already exists; use an empty output directory") + os.symlink(entry.resolve(), target) + + +def sha256_file(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1 << 24), b""): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--base", required=True, help="MiniMax-H3 snapshot at the adapter's pinned revision") + parser.add_argument("--adapter", required=True, help="adapter file, release directory, or Hub repo id") + parser.add_argument("--task", required=True, choices=("fl2va", "ref2va")) + parser.add_argument("--revision", default=None, help="Hub revision when --adapter is a repo id") + parser.add_argument("--output", required=True) + args = parser.parse_args() + + adapter = load_adapter(args.adapter, task=args.task, revision=args.revision) + digest = sha256_file(adapter.path) + if adapter.manifest.get("sha256") and digest != adapter.manifest["sha256"]: + raise ValueError(f"{adapter.path} has sha256 {digest}, the manifest says {adapter.manifest['sha256']}") + base, output = Path(args.base).resolve(), Path(args.output) + if output.exists() and any(output.iterdir()): + raise FileExistsError(f"{output} is not empty") + factors = adapter_factors(adapter.path) + scale = float(adapter.manifest["alpha"]) / float(adapter.manifest["rank"]) + link_base(base, output, adapter.trunk) + report = merge_trunk(base / adapter.trunk, output / adapter.trunk, factors, scale) + contract = { + "schema": CONTRACT_SCHEMA, + "task": args.task, + "trunk": adapter.trunk, + "release_id": adapter.manifest.get("release_id"), + "adapter_sha256": digest, + "base": adapter.manifest.get("base"), + "numerics_policy": adapter.manifest.get("numerics_policy"), + "denoising_steps": list(PIKA_DENOISING_STEPS), + "video_scheduler_shift": PIKA_VIDEO_SCHEDULER_SHIFT, + "audio_scheduler_shift": PIKA_AUDIO_SCHEDULER_SHIFT, + "sampler": "renoise", + "merged_modules": report["modules"], + "merge": "fp32 W + (alpha/rank) B@A, rounded once to the base dtype", + } + (output / CONTRACT_FILE).write_text(json.dumps(contract, indent=2) + "\n") + (output / "merge_report.json").write_text(json.dumps(report, indent=2) + "\n") + print(f"merged {report['modules']} modules into {output / adapter.trunk} " + f"(max relative weight change {report['max_rel_delta']:.4f})") + + +if __name__ == "__main__": + main()