[feat]: Pika H3 Flex NVFP4 QAD profile (Ref2VA + FL2VA, six-forward re-noise ladder)
Add MiniMaxH3FlexRef2VAModel / MiniMaxH3FlexFL2VAModel for OmniRefPDDQADMethod: dense attention, Pika's shifted 1000..40 ladder, re-noise steps whose noise is a function of (row seed, forward, modality), NVFP4 linears only. Add an offline fp32 adapter merge script and QAD configs for both adapters. Let model profiles pick their attention backend and receive the prepared row in schedulers().
This commit is contained in:
Aryan Kumar
@@ -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: {}
|
||||
@@ -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: {}
|
||||
@@ -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))
|
||||
@@ -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):
|
||||
|
||||
@@ -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"
|
||||
]
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 <MiniMax-H3 snapshot @42ed227> --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()
|
||||
Reference in New Issue
Block a user