diff --git a/docs-site/lib/cookbook-pages.ts b/docs-site/lib/cookbook-pages.ts index 3dac83a84..5c3d6b4ed 100644 --- a/docs-site/lib/cookbook-pages.ts +++ b/docs-site/lib/cookbook-pages.ts @@ -48,6 +48,15 @@ export const recipePages = [ tags: ["Reference to video", "Audio", "NVFP4"], featured: true, }, + { + id: "open-h3-regenerate-2k", + family: "MiniMax H3", + title: "Open-H3-Regenerate-2K", + mark: "H3", + path: "/cookbook/minimax-h3/open-h3-regenerate-2k", + description: "Regenerate a 768P MiniMax-H3 draft at 2560x1440 with a three-step refiner.", + tags: ["Video to video", "2K", "3-step"], + }, { id: "fastwan21", family: "Wan", diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index e19feae8a..4eac9fe00 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -760,6 +760,22 @@ ], "limitations": ["Pass --reference-audio for an additional audio reference. The checked-in source defaults to four GPUs."] }, + { + "id": "open-h3-regenerate-2k", + "family": "minimax_h3", + "stage": "inference", + "task": "Video to 2K video", + "label": "Open-H3-Regenerate-2K", + "summary": "Regenerate a 768P MiniMax-H3 draft at 2560x1440 with NVlabs' three-step refiner. The output is silent.", + "model": "Efficient-Large-Model/Open-H3-Regenerate-2K", + "source": "examples/inference/basic/basic_open_h3_regenerate_2k.py", + "command": ": \"${INPUT_VIDEO:?Set INPUT_VIDEO to a 768P MiniMax-H3 draft}\"\npython examples/inference/basic/basic_open_h3_regenerate_2k.py --video \"$INPUT_VIDEO\" --prompt \"A red fox trots through fresh snow in a pine forest at golden hour.\"", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "2560x1440 MP4 at outputs/open_h3_regenerate_2k/regen2k_0.mp4", + "related": ["minimax-h3-t2v", "fasth3-8step-v2-cuda"] + }, { "id": "fasth3-lora-preview", "family": "minimax_h3", diff --git a/docs/cookbook/options.yaml b/docs/cookbook/options.yaml index 627ac1e79..6953f915a 100644 --- a/docs/cookbook/options.yaml +++ b/docs/cookbook/options.yaml @@ -44,7 +44,7 @@ runtimes: host: 0.0.0.0 port: 8000 served_model_name: null - workloads: [t2v, i2v] + workloads: [t2v, i2v, v2v] options: &cuda_options server.served_model_name: anyOf: diff --git a/docs/cookbook/recipes/open-h3-regenerate-2k.yaml b/docs/cookbook/recipes/open-h3-regenerate-2k.yaml new file mode 100644 index 000000000..b3f3aaf69 --- /dev/null +++ b/docs/cookbook/recipes/open-h3-regenerate-2k.yaml @@ -0,0 +1,18 @@ +title: Open-H3-Regenerate-2K +summary: Regenerate a 768P MiniMax-H3 draft at 2560x1440 with three refine steps. + +deployments: + cuda-rest: + runtime: fastvideo-cuda-rest + guide: docs/inference/open_h3_regenerate_2k.md + workload: v2v + defaults: examples/serving/openai_open_h3_regenerate_2k.yaml + requirements: + - 'UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]"' + - "One GPU with about 150 GB free (B200 or GB200). Send the draft as video_reference; local paths must sit under server.media_roots." + - "Optional NVFP4 decoder (Blackwell only): hf download FastVideo/FastH3-NVFP4-VAE-Decoder --include 'quality-36L/decoder-nvfp4.safetensors' --local-dir fasth3-nvfp4-vae" + options: + - server.* + - generator.engine.offload.text_encoder + - generator.pipeline.experimental.h3_vae_nvfp4_decoder + - default_request.sampling.seed diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 9027bb660..b1353dac9 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -153,6 +153,12 @@ surfaces: pipeline_config_extensions: preset_owned: + draft_size: + sources: + - fastvideo.configs.pipelines.open_h3_regenerate_2k.OpenH3Regenerate2KPipelineConfig + refine_start_sigma: + sources: + - fastvideo.configs.pipelines.open_h3_regenerate_2k.OpenH3Regenerate2KPipelineConfig audio_encoder_precision: sources: - fastvideo.configs.pipelines.wan.WanS2V14BConfig diff --git a/docs/inference/open_h3_regenerate_2k.md b/docs/inference/open_h3_regenerate_2k.md index 6081e57ae..b7beffbb9 100644 --- a/docs/inference/open_h3_regenerate_2k.md +++ b/docs/inference/open_h3_regenerate_2k.md @@ -87,6 +87,8 @@ The response headers carry `X-Request-Id`, `X-Model`, `X-Inference-Time-S` and ` - `size`, `width` and `height` may be omitted or `2560x1440`; `num_frames` may be omitted or 124 (the frame count of the VAE grid; the video has 121 frames); `fps` must be 24. `seconds` is not accepted. - `num_inference_steps` is 3 (default) or 2; `seed` sets the refine noise. +- Serve with `attention_backend: TORCH_SDPA`: the refine stage selects cuDNN attention, which at about 133k tokens is + about 3x faster than PyTorch's default flash dispatch. The `FLASH_ATTN` backend bypasses that selection. ## NVFP4 video decoder @@ -153,7 +155,7 @@ keeping structure; use it only when the decode budget is tight. ## Parity with the reference -Measured on one GB200 (tray hpc-rack-1-5, BF16, dense attention) against the reference `standard_fwd` with the released +Measured on one GB200 (BF16, `TORCH_SDPA` with cuDNN attention) against the reference `standard_fwd` with the released VAE, same draft, seed (1766580932), prompt and VAE posterior seed (42). The reference's noise is a CUDA generator draw that FastVideo reproduces (the same `torch.Generator(device="cuda")` seed), so no injected noise was needed. diff --git a/docs/js/cookbook-workloads.mjs b/docs/js/cookbook-workloads.mjs index d2820a60a..243ee687c 100644 --- a/docs/js/cookbook-workloads.mjs +++ b/docs/js/cookbook-workloads.mjs @@ -32,6 +32,14 @@ export const WORKLOADS = Object.freeze({ extension: "wav", media_type: "audio/wav", }), + v2v: Object.freeze({ + input: "video", + output: "video", + endpoint: "/v1/videos/sync", + extension: "mp4", + media_type: "video/mp4", + reference: "video_reference", + }), v2a: Object.freeze({ input: "video", output: "audio", diff --git a/docs/tests/cookbook-config.test.mjs b/docs/tests/cookbook-config.test.mjs index 36c26869f..d95640fe5 100644 --- a/docs/tests/cookbook-config.test.mjs +++ b/docs/tests/cookbook-config.test.mjs @@ -488,6 +488,7 @@ test("sample requests follow shared workload media", () => { t2i: { endpoint: /\/v1\/images\/generations/, output: /--output output\.json/, reference: null }, t2a: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.wav/, reference: null }, v2a: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.wav/, reference: "video_reference" }, + v2v: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.mp4/, reference: "video_reference" }, }; for (const [workload, expected] of Object.entries(expectations)) { const data = catalog(); diff --git a/examples/inference/basic/basic_open_h3_regenerate_2k.py b/examples/inference/basic/basic_open_h3_regenerate_2k.py index e2af9c59e..262a371a4 100644 --- a/examples/inference/basic/basic_open_h3_regenerate_2k.py +++ b/examples/inference/basic/basic_open_h3_regenerate_2k.py @@ -38,7 +38,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--output", type=Path, default=Path("outputs/open_h3_regenerate_2k")) parser.add_argument("--steps", type=int, choices=(2, 3), default=3, help="refine steps (3 is the released setting)") parser.add_argument("--seed", type=int, default=1766580932) - parser.add_argument("--attention-backend", default="FLASH_ATTN") + parser.add_argument("--attention-backend", default="TORCH_SDPA") parser.add_argument("--draft-size", choices=("768p", "576p"), default="768p", help="size the draft is encoded at; 576p for drafts from other video models") parser.add_argument("--vae-decoder", default=None, diff --git a/examples/serving/openai_open_h3_regenerate_2k.yaml b/examples/serving/openai_open_h3_regenerate_2k.yaml index 225bc79e5..6744ab17b 100644 --- a/examples/serving/openai_open_h3_regenerate_2k.yaml +++ b/examples/serving/openai_open_h3_regenerate_2k.yaml @@ -17,7 +17,7 @@ generator: vae_enabled: false pipeline: experimental: - attention_backend: FLASH_ATTN + attention_backend: TORCH_SDPA # 768p for MiniMax-H3 drafts; 576p encodes drafts from other video models at 1024x576. draft_size: 768p # NVFP4 video decoder (Blackwell only). Download one tier from FastVideo/FastH3-NVFP4-VAE-Decoder and point diff --git a/fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py b/fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py index cd3e605ab..74c0ca6a0 100644 --- a/fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py +++ b/fastvideo-kernel/benchmarks/bench_vsa_tile_layout.py @@ -5,7 +5,7 @@ Benchmark the SM100 fused VSA64 tile layout against the legacy scatter/transpose Times Wan VSA inference preprocessing through the backend (``VideoSparseAttentionImpl.preprocess_qkv``), toggling ``FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT``: - - legacy: tile scatter into BSHD + four BSHD->BHSD transposes + - legacy: tile scatter into BSHD + three contiguous Q/K/V transposes + a gate view - fused: one Triton kernel writing tiled BHSD directly With ``--include_attn`` the timed region also covers ``forward`` (the @@ -59,20 +59,20 @@ def parse_arguments() -> argparse.Namespace: def make_step(impl: VideoSparseAttentionImpl, qkvg: torch.Tensor, builder: VideoSparseAttentionMetadataBuilder, - build_kwargs: dict, fused: bool, include_attn: bool) -> Callable[[], torch.Tensor]: + build_kwargs: dict, fused: bool, include_attn: bool) -> Callable[[], tuple[torch.Tensor, ...]]: metadata = builder.build(**build_kwargs) - def step() -> torch.Tensor: + def step() -> tuple[torch.Tensor, ...]: # The routing gate reads the switch on every call. with envs.FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT.override(not fused): tiled = impl.preprocess_qkv(qkvg, metadata) q, k, v, gate = tiled.chunk(4) if include_attn: - return impl.forward(q, k, v, gate, metadata) + return (impl.forward(q, k, v, gate, metadata), ) if metadata.fused_layout_active: - return tiled - # Legacy route: the transposes forward() applies before video_sparse_attn. - return torch.cat([t.transpose(1, 2).contiguous() for t in (q, k, v, gate)]) + return q, k, v, gate + return (q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(), + gate.transpose(1, 2)) with torch.no_grad(): step() @@ -100,7 +100,7 @@ def main() -> None: legacy = make_step(impl, qkvg, builder, build_kwargs, fused=False, include_attn=args.include_attn) fused = make_step(impl, qkvg, builder, build_kwargs, fused=True, include_attn=args.include_attn) with torch.no_grad(): - if not torch.equal(legacy(), fused()): + if not all(torch.equal(old, new) for old, new in zip(legacy(), fused(), strict=True)): raise RuntimeError("Fused and legacy outputs differ.") legacy_ms = do_bench(legacy, warmup=args.warmup, rep=args.rep, return_mode="median") fused_ms = do_bench(fused, warmup=args.warmup, rep=args.rep, return_mode="median") diff --git a/fastvideo-kernel/tests/test_vsa_tile_layout.py b/fastvideo-kernel/tests/test_vsa_tile_layout.py index 13988232c..b02dd96ca 100644 --- a/fastvideo-kernel/tests/test_vsa_tile_layout.py +++ b/fastvideo-kernel/tests/test_vsa_tile_layout.py @@ -5,6 +5,9 @@ Production only routes SM100 here, but the kernel is plain Triton, so parity is checked on any CUDA device. """ +import importlib.util +from pathlib import Path + import pytest import torch @@ -40,6 +43,63 @@ def test_tile_to_bhsd_matches_scatter_and_four_transposes(dit_shape): source_index[metadata.non_pad_index] = metadata.tile_partition_indices.to(torch.int32) out = torch.empty((4, 2, padded_sequence, 128), device="cuda", dtype=torch.bfloat16) actual = tile_to_bhsd(qkvg, source_index, out, sequence, padded_sequence, 2, 128).chunk(4) - for name, old, new in zip(("q", "k", "v", "gate"), expected, actual): + for name, old, new in zip(("q", "k", "v", "gate"), expected, actual, strict=True): assert old.shape == new.shape, name assert torch.equal(old, new), name + + +@pytest.mark.parametrize("fused", [False, True]) +def test_layout_benchmark_matches_forward_preprocessing(monkeypatch, fused): + if not torch.cuda.is_available(): + pytest.skip("CUDA device required") + + from fastvideo.attention.backends import video_sparse_attn as vsa + + path = Path(__file__).resolve().parents[1] / "benchmarks/bench_vsa_tile_layout.py" + spec = importlib.util.spec_from_file_location("bench_vsa_tile_layout", path) + benchmark = importlib.util.module_from_spec(spec) + spec.loader.exec_module(benchmark) + original_gate = vsa._can_fuse_vsa64_layout + + def allow_test_device(x, metadata): + with monkeypatch.context() as gate_patch: + gate_patch.setattr(torch.cuda, "get_device_capability", lambda *args: (10, 0)) + return original_gate(x, metadata) + + monkeypatch.setattr(vsa, "_can_fuse_vsa64_layout", allow_test_device) + impl = object.__new__(vsa.VideoSparseAttentionImpl) + builder = vsa.VideoSparseAttentionMetadataBuilder() + kwargs = dict(current_timestep=0, + raw_latent_shape=(5, 5, 6), + patch_size=(1, 1, 1), + VSA_sparsity=0.8, + device=torch.device("cuda"), + cache_tile_buf=True) + qkvg = torch.randn((4, 157, 2, 128), device="cuda", dtype=torch.bfloat16)[:, :150] + captured = [] + + def capture_attention_inputs(q, k, v, *args, **kwargs): + captured[:] = [q, k, v, kwargs["compress_attn_weight"]] + return q + + monkeypatch.setattr(vsa, "video_sparse_attn", capture_attention_inputs) + with torch.no_grad(): + layout_step = benchmark.make_step(impl, qkvg, builder, kwargs, fused=fused, include_attn=False) + attention_step = benchmark.make_step(impl, qkvg, builder, kwargs, fused=fused, include_attn=True) + + def reject_cat(*args, **kwargs): + pytest.fail("The timed layout step must not concatenate its outputs") + + with monkeypatch.context() as timed_patch: + timed_patch.setattr(torch, "cat", reject_cat) + layout = layout_step() + attention = attention_step() + + assert isinstance(layout, tuple) and len(layout) == 4 + for actual, expected in zip(layout, captured, strict=True): + assert torch.equal(actual, expected) + assert actual.stride() == expected.stride() + assert all(t.is_contiguous() for t in layout[:3]) + assert layout[3].is_contiguous() is fused + assert isinstance(attention, tuple) and len(attention) == 1 + assert torch.equal(attention[0], captured[0].transpose(1, 2)) diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index 1458ca52b..31b21d531 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -254,13 +254,7 @@ class BaseLayerWithLoRA(nn.Module): current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()).full_tensor() - # Apply LoRA with alpha scaling - lora_delta = ( - self.slice_lora_b_weights(self.lora_B).to(data) @ self.slice_lora_a_weights(self.lora_A).to(data)) - scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0) - lora_delta *= scale * self.lora_strength - data += lora_delta - unsharded_base_layer.weight = nn.Parameter(data.to(current_device)) + unsharded_base_layer.weight = nn.Parameter(self._merged_weight(data).to(current_device)) if isinstance(getattr(self.base_layer, "bias", None), DTensor): unsharded_base_layer.bias = nn.Parameter( self.base_layer.bias.to(get_local_torch_device(), @@ -277,18 +271,22 @@ class BaseLayerWithLoRA(nn.Module): current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()) - # Apply LoRA with alpha scaling. An adapter stored wider than the base (an FP32 adapter on a BF16 DiT) - # is multiplied and added at its own precision and rounded once, as a merge-then-cast reference does. - accumulate_dtype = torch.promote_types(data.dtype, self.lora_A.dtype) - lora_delta = (self.slice_lora_b_weights(self.lora_B.to(accumulate_dtype)) @ self.slice_lora_a_weights( - self.lora_A.to(accumulate_dtype))) - scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0) - lora_delta *= scale * self.lora_strength - data = (data.to(accumulate_dtype) + lora_delta).to(data.dtype) - self.base_layer.weight.data = data.to(current_device, non_blocking=True) + self.base_layer.weight.data = self._merged_weight(data).to(current_device, non_blocking=True) self.merged = True + def _merged_weight(self, weight: torch.Tensor) -> torch.Tensor: + """``weight`` plus the scaled adapter delta, in ``weight``'s dtype. + + An adapter stored wider than the base (an FP32 adapter on a BF16 DiT) is multiplied and added at its own + precision and rounded once, as a merge-then-cast reference does; with FSDP or without. + """ + accumulate_dtype = torch.promote_types(weight.dtype, self.lora_A.dtype) + lora_B = self.slice_lora_b_weights(self.lora_B.to(device=weight.device, dtype=accumulate_dtype)) + lora_A = self.slice_lora_a_weights(self.lora_A.to(device=weight.device, dtype=accumulate_dtype)) + scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0) + return (weight.to(accumulate_dtype) + (lora_B @ lora_A) * (scale * self.lora_strength)).to(weight.dtype) + @torch.no_grad() # @torch.compile(dynamic=True) def unmerge_lora_weights(self, *, if_merged_only: bool = True) -> None: diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index a9e3b801a..9c9ffb90f 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -1655,7 +1655,6 @@ class PipelineComponentLoader: component_model_path, ) - # A component under a subfolder (e.g. "base_model/vae") dispatches on its own name. module_name = module_name.rsplit("/", 1)[-1] # Get the appropriate loader for this module type loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers) diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py index b5bf46360..2a5868094 100644 --- a/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py @@ -54,7 +54,6 @@ class OpenH3Regenerate2KPipeline(MiniMaxH3Pipeline): revision=revision, allow_patterns=self.get_hf_download_allow_patterns()) if self.fastvideo_args.lora_path is None: - # The merged refine LoRA is part of the release; it is a constructor-time adapter like any H3 LoRA. self.fastvideo_args.lora_path = os.path.join(self.model_path, LORA_FILE) return verify_model_config_and_directory(os.path.join(self.model_path, BASE_MODEL_DIR), required_component_dirs=()) diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py index dc0645702..a08efab22 100644 --- a/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py @@ -34,6 +34,9 @@ def interpolate_curve(table: torch.Tensor, timesteps: torch.Tensor) -> torch.Ten return torch.lerp(table[low], table[low + 1], (position - low).unsqueeze(-1)) +_ATTACHED_ATTR = "_open_h3_sharp_curve_branches" + + class SharpCurveBranches: def __init__(self, weights_dir: str | Path) -> None: @@ -50,15 +53,17 @@ class SharpCurveBranches: if len(self.factors) != NUM_CURVE_BRANCHES: raise ValueError(f"Expected {NUM_CURVE_BRANCHES} Sharp curve branches, found {len(self.factors)}.") self.coords: torch.Tensor | None = None - self._attached: set[int] = set() self._device_factors: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} def set_timesteps(self, unique_timesteps: torch.Tensor) -> None: self.coords = interpolate_curve(self.table, unique_timesteps) def attach(self, transformer: torch.nn.Module) -> None: - """Register the branch hooks on ``transformer`` once.""" - if id(transformer) in self._attached: + """Register the branch hooks on ``transformer`` unless this instance already carries them. + + The mark lives on the module, so a reloaded lazy transformer, a new instance, gets its own hooks. + """ + if getattr(transformer, _ATTACHED_ATTR, None) is self: return modules = dict(transformer.named_modules()) missing = sorted(set(self.factors) - set(modules)) @@ -66,7 +71,7 @@ class SharpCurveBranches: raise ValueError(f"Sharp curve branches match no transformer module: {missing[:5]}") for name in self.factors: modules[name].register_forward_hook(self._hook(name)) - self._attached.add(id(transformer)) + setattr(transformer, _ATTACHED_ATTR, self) def _hook(self, name: str): diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py index 6dba2d58f..4ffef9cab 100644 --- a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py @@ -53,9 +53,7 @@ class Regenerate2KDenoisingStage(PipelineStage): @torch.no_grad() def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: transformer = self.transformer - if is_lazy_module(transformer): - transformer.materialize() - self.curve.attach(transformer) + self.curve.attach(transformer.materialize() if is_lazy_module(transformer) else transformer) layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY) if not isinstance(layout, MiniMaxH3PackedLayout): raise ValueError("The MiniMax-H3 packed layout is missing before denoising.") diff --git a/fastvideo/tests/layers/test_lora_merge_precision.py b/fastvideo/tests/layers/test_lora_merge_precision.py index dea0e3442..6bbf39ca0 100644 --- a/fastvideo/tests/layers/test_lora_merge_precision.py +++ b/fastvideo/tests/layers/test_lora_merge_precision.py @@ -42,3 +42,20 @@ def test_fp32_adapter_is_added_in_fp32_and_rounded_once() -> None: def test_bf16_adapter_keeps_the_bf16_merge() -> None: merged, base, a, b = _merged(torch.bfloat16) assert torch.equal(merged, base + b @ a) + + +def test_merge_helper_used_with_and_without_fsdp_rounds_once() -> None: + torch.manual_seed(1) + layer = ReplicatedLinear(IN, OUT, bias=False, params_dtype=torch.bfloat16) + with torch.no_grad(): + layer.weight.copy_(torch.randn(OUT, IN)) + wrapped = get_lora_layer(layer) + a, b = torch.randn(RANK, IN), torch.randn(OUT, RANK) * 0.1 + wrapped.lora_A, wrapped.lora_B = torch.nn.Parameter(a), torch.nn.Parameter(b) + wrapped.lora_rank, wrapped.lora_alpha, wrapped.lora_strength = RANK, RANK, 0.5 + weight = layer.weight.detach().clone() + + merged = wrapped._merged_weight(weight) + + assert merged.dtype == torch.bfloat16 + assert torch.equal(merged, (weight.float() + 0.5 * (b @ a)).to(torch.bfloat16)) diff --git a/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py index f843094ad..0b60d57d2 100644 --- a/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py +++ b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py @@ -124,7 +124,7 @@ def test_sharp_branch_adds_curve_term_to_projection_output(tmp_path): block.proj = _Projection(6) model.blocks.append(block) curve.attach(model) - curve.attach(model) # idempotent: a second attach must not double the term + curve.attach(model) timesteps = torch.tensor([0.1, 0.9]) x = torch.randn(2, 4) base, _ = model.blocks[0].proj.linear(x), None @@ -199,15 +199,51 @@ class _RecordingCurve: def __init__(self): self.timesteps = [] self.attached = 0 + self.targets = [] def attach(self, transformer): self.attached += 1 + self.targets.append(transformer) def set_timesteps(self, unique): self.timesteps.append(unique.clone()) -def _denoise(monkeypatch, steps, velocity=0.5): +def _curve_model(out_features): + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList() + for _ in range(NUM_CURVE_BRANCHES): + block = torch.nn.Module() + block.proj = _Projection(out_features) + model.blocks.append(block) + return model + + +def test_sharp_branches_reach_a_reloaded_transformer_instance(tmp_path): + _write_curve_release(tmp_path, 6, NUM_CURVE_BRANCHES) + curve = SharpCurveBranches(tmp_path) + curve.set_timesteps(torch.tensor([0.1, 0.9])) + x = torch.randn(2, 4) + for model in (_curve_model(6), _curve_model(6)): + curve.attach(model) + out, _ = model.blocks[0].proj(x) + assert not torch.allclose(out, model.blocks[0].proj.linear(x)) + + +def test_stage_attaches_the_branches_to_each_materialized_lazy_transformer(monkeypatch): + from fastvideo.pipelines.lazy_module import LazyModule + + lazy = LazyModule("transformer", lambda: _ConstantVelocityTransformer(0.5)) + first = _denoise(monkeypatch, steps=2, transformer=lazy)[2] + loaded_first = first.targets[-1] + assert lazy.release() + second = _denoise(monkeypatch, steps=2, transformer=lazy, curve=first)[2] + + assert not isinstance(loaded_first, LazyModule) and not isinstance(second.targets[-1], LazyModule) + assert second.targets[-1] is not loaded_first + + +def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None): from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import denoising @@ -221,7 +257,8 @@ def _denoise(monkeypatch, steps, velocity=0.5): batch.latents = torch.zeros(layout.video_indices.numel(), 8) batch.audio_latents = torch.ones(layout.audio_indices.numel(), 4) batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout - transformer, curve = _ConstantVelocityTransformer(velocity), _RecordingCurve() + transformer = transformer if transformer is not None else _ConstantVelocityTransformer(velocity) + curve = curve if curve is not None else _RecordingCurve() stage = denoising.Regenerate2KDenoisingStage(transformer, curve, None) return stage.forward(batch, None), transformer, curve