diff --git a/fastvideo/pipelines/basic/minimax_h3/disaggregated.py b/fastvideo/pipelines/basic/minimax_h3/disaggregated.py index 059c86213..4c85f45e6 100644 --- a/fastvideo/pipelines/basic/minimax_h3/disaggregated.py +++ b/fastvideo/pipelines/basic/minimax_h3/disaggregated.py @@ -8,7 +8,6 @@ from uuid import uuid4 import torch -from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import ( MiniMaxH3BasePipeline, @@ -28,23 +27,23 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, PipelineLoggin MINIMAX_H3_WIRE_SCHEMA_VERSION = 1 -def _cpu_tensor(value: torch.Tensor, name: str, dimensions: int) -> torch.Tensor: +def _cuda_tensor(value: torch.Tensor, name: str, dimensions: int) -> torch.Tensor: if not isinstance(value, torch.Tensor) or value.ndim != dimensions: shape = None if not isinstance(value, torch.Tensor) else tuple(value.shape) raise ValueError(f"MiniMax-H3 wire field {name!r} must be a {dimensions}D tensor, got {shape}.") - return value.detach().to(device="cpu").contiguous() + return value.detach().to(device="cuda", non_blocking=True).contiguous() -def _cpu_layout(layout: MiniMaxH3PackedLayout) -> MiniMaxH3PackedLayout: +def _cuda_layout(layout: MiniMaxH3PackedLayout) -> MiniMaxH3PackedLayout: if not isinstance(layout, MiniMaxH3PackedLayout): raise TypeError("MiniMax-H3 encode output is missing its packed layout.") return replace( layout, - position_ids=_cpu_tensor(layout.position_ids, "layout.position_ids", 2), - token_tags=_cpu_tensor(layout.token_tags, "layout.token_tags", 1), - video_indices=_cpu_tensor(layout.video_indices, "layout.video_indices", 1), - audio_indices=_cpu_tensor(layout.audio_indices, "layout.audio_indices", 1), - text_indices=_cpu_tensor(layout.text_indices, "layout.text_indices", 1), + position_ids=_cuda_tensor(layout.position_ids, "layout.position_ids", 2), + token_tags=_cuda_tensor(layout.token_tags, "layout.token_tags", 1), + video_indices=_cuda_tensor(layout.video_indices, "layout.video_indices", 1), + audio_indices=_cuda_tensor(layout.audio_indices, "layout.audio_indices", 1), + text_indices=_cuda_tensor(layout.text_indices, "layout.text_indices", 1), ) @@ -61,7 +60,7 @@ def _raw_latent_shape(batch: ForwardBatch, stage: str) -> tuple[int, int, int, i return (int(shape[0]), int(shape[1]), int(shape[2]), int(shape[3]), int(shape[4])) -def _validate_cpu_layout(layout: MiniMaxH3PackedLayout) -> None: +def _validate_cuda_layout(layout: MiniMaxH3PackedLayout) -> None: if not isinstance(layout, MiniMaxH3PackedLayout): raise TypeError("MiniMax-H3 wire payload requires a packed layout.") expected_shapes = { @@ -75,8 +74,8 @@ def _validate_cpu_layout(layout: MiniMaxH3PackedLayout) -> None: ("audio_indices", layout.audio_indices), ("text_indices", layout.text_indices), ): - if tensor.device.type != "cpu" or not tensor.is_contiguous(): - raise ValueError(f"MiniMax-H3 wire field layout.{name} must be a contiguous CPU tensor.") + if not tensor.is_cuda or not tensor.is_contiguous(): + raise ValueError(f"MiniMax-H3 wire field layout.{name} must be a contiguous CUDA tensor.") expected = expected_shapes.get(name) if expected is not None and tuple(tensor.shape) != expected: raise ValueError( @@ -87,7 +86,7 @@ def _validate_cpu_layout(layout: MiniMaxH3PackedLayout) -> None: @dataclass(frozen=True) class MiniMaxH3EncodedState: - """Minimal CPU payload sent from the encoder/VAE node to the DiT node.""" + """Minimal CUDA payload sent from the encoder/VAE node to the DiT node.""" request_id: str prompt_embeds: torch.Tensor @@ -118,9 +117,9 @@ class MiniMaxH3EncodedState: ("video_latents", self.video_latents, 2), ("audio_latents", self.audio_latents, 2), ): - if tensor.device.type != "cpu" or not tensor.is_contiguous() or tensor.ndim != dimensions: - raise ValueError(f"MiniMax-H3 wire field {name!r} must be a contiguous CPU {dimensions}D tensor.") - _validate_cpu_layout(self.layout) + if not tensor.is_cuda or not tensor.is_contiguous() or tensor.ndim != dimensions: + raise ValueError(f"MiniMax-H3 wire field {name!r} must be a contiguous CUDA {dimensions}D tensor.") + _validate_cuda_layout(self.layout) if self.video_latents.shape[0] != self.layout.video_indices.numel(): raise ValueError("MiniMax-H3 video latent rows do not match layout.video_indices.") if self.audio_latents.shape[0] != self.layout.audio_indices.numel(): @@ -135,10 +134,10 @@ class MiniMaxH3EncodedState: layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY) return cls( request_id=request_id if request_id is not None else uuid4().hex, - prompt_embeds=_cpu_tensor(batch.prompt_embeds[0], "prompt_embeds", 3), - video_latents=_cpu_tensor(batch.latents, "video_latents", 2), - audio_latents=_cpu_tensor(batch.audio_latents, "audio_latents", 2), - layout=_cpu_layout(layout), + prompt_embeds=_cuda_tensor(batch.prompt_embeds[0], "prompt_embeds", 3), + video_latents=_cuda_tensor(batch.latents, "video_latents", 2), + audio_latents=_cuda_tensor(batch.audio_latents, "audio_latents", 2), + layout=_cuda_layout(layout), raw_latent_shape=_raw_latent_shape(batch, "encode"), num_inference_steps=int(batch.num_inference_steps), vsa_sparsity=float(batch.VSA_sparsity), @@ -148,16 +147,15 @@ class MiniMaxH3EncodedState: logging_info=batch.logging_info, ) - def to_batch(self, *, device: torch.device | str | None = None) -> ForwardBatch: + def to_batch(self) -> ForwardBatch: # Pickle/Ray reconstruction does not invoke dataclass __post_init__. - # Revalidate the schema and tensor contract at the receiving boundary. + # Revalidate the schema and CUDA tensor contract at the receiving boundary. self.__post_init__() - target = get_local_torch_device() if device is None else torch.device(device) return ForwardBatch( data_type="video", - prompt_embeds=[self.prompt_embeds.to(target)], - latents=self.video_latents.to(target), - audio_latents=self.audio_latents.to(target), + prompt_embeds=[self.prompt_embeds], + latents=self.video_latents, + audio_latents=self.audio_latents, raw_latent_shape=self.raw_latent_shape, num_inference_steps=self.num_inference_steps, VSA_sparsity=self.vsa_sparsity, @@ -174,7 +172,7 @@ class MiniMaxH3EncodedState: @dataclass(frozen=True) class MiniMaxH3DenoisedState: - """Minimal CPU payload returned by the DiT node for VAE decoding.""" + """Minimal CUDA payload returned by the DiT node for VAE decoding.""" request_id: str video_latents: torch.Tensor @@ -191,9 +189,9 @@ class MiniMaxH3DenoisedState: if len(self.raw_latent_shape) != 5 or min(self.raw_latent_shape) <= 0: raise ValueError(f"Invalid MiniMax-H3 raw latent shape: {self.raw_latent_shape}.") for name, tensor in (("video_latents", self.video_latents), ("audio_latents", self.audio_latents)): - if tensor.device.type != "cpu" or not tensor.is_contiguous() or tensor.ndim != 2: - raise ValueError(f"MiniMax-H3 wire field {name!r} must be a contiguous CPU 2D tensor.") - _validate_cpu_layout(self.layout) + if not tensor.is_cuda or not tensor.is_contiguous() or tensor.ndim != 2: + raise ValueError(f"MiniMax-H3 wire field {name!r} must be a contiguous CUDA 2D tensor.") + _validate_cuda_layout(self.layout) if self.video_latents.shape[0] != self.layout.video_indices.numel(): raise ValueError("MiniMax-H3 video latent rows do not match layout.video_indices.") if self.audio_latents.shape[0] != self.layout.audio_indices.numel(): @@ -205,9 +203,9 @@ class MiniMaxH3DenoisedState: raise ValueError("MiniMax-H3 denoise must return both latent streams.") return cls( request_id=request_id, - video_latents=_cpu_tensor(batch.latents, "video_latents", 2), - audio_latents=_cpu_tensor(batch.audio_latents, "audio_latents", 2), - layout=_cpu_layout(batch.extra.get(MINIMAX_H3_LAYOUT_KEY)), + video_latents=_cuda_tensor(batch.latents, "video_latents", 2), + audio_latents=_cuda_tensor(batch.audio_latents, "audio_latents", 2), + layout=_cuda_layout(batch.extra.get(MINIMAX_H3_LAYOUT_KEY)), raw_latent_shape=_raw_latent_shape(batch, "denoise"), logging_info=batch.logging_info, ) diff --git a/fastvideo/tests/worker/test_minimax_h3_disaggregated.py b/fastvideo/tests/worker/test_minimax_h3_disaggregated.py index 248b5a836..08ff48d43 100644 --- a/fastvideo/tests/worker/test_minimax_h3_disaggregated.py +++ b/fastvideo/tests/worker/test_minimax_h3_disaggregated.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""CPU contracts for the two-node MiniMax-H3 component pipeline.""" +"""CUDA contracts for the two-node MiniMax-H3 component pipeline.""" from __future__ import annotations @@ -38,16 +38,23 @@ from fastvideo.worker.executor import Executor def _layout(*, noncontiguous: bool = False) -> MiniMaxH3PackedLayout: - position_ids = torch.arange(27, dtype=torch.float64).reshape(3, 9).transpose(0, 1) + position_ids = ( + torch.arange(27, dtype=torch.float64, device="cuda") + .reshape(3, 9) + .transpose(0, 1) + ) if not noncontiguous: position_ids = position_ids.contiguous() return MiniMaxH3PackedLayout( sequence_length=9, position_ids=position_ids, - token_tags=torch.tensor([1, 1, 1, 2, 2, 0, 0, 0, 0]), - video_indices=torch.tensor([5, 6, 7, 8]), - audio_indices=torch.tensor([3, 4]), - text_indices=torch.tensor([0, 1, 2]), + token_tags=torch.tensor( + [1, 1, 1, 2, 2, 0, 0, 0, 0], + device="cuda", + ), + video_indices=torch.tensor([5, 6, 7, 8], device="cuda"), + audio_indices=torch.tensor([3, 4], device="cuda"), + text_indices=torch.tensor([0, 1, 2], device="cuda"), num_condition_video_rows=0, num_condition_audio_rows=0, num_video_latent_frames=2, @@ -58,9 +65,21 @@ def _layout(*, noncontiguous: bool = False) -> MiniMaxH3PackedLayout: def _encoded_batch(*, request_id: str | None = None) -> ForwardBatch: - prompt_embeds = torch.arange(12, dtype=torch.float32).reshape(1, 4, 3).transpose(1, 2) - video_latents = torch.arange(20, dtype=torch.float32).reshape(5, 4).transpose(0, 1) - audio_latents = torch.arange(12, dtype=torch.float32).reshape(6, 2).transpose(0, 1) + prompt_embeds = ( + torch.arange(12, dtype=torch.float32, device="cuda") + .reshape(1, 4, 3) + .transpose(1, 2) + ) + video_latents = ( + torch.arange(20, dtype=torch.float32, device="cuda") + .reshape(5, 4) + .transpose(0, 1) + ) + audio_latents = ( + torch.arange(12, dtype=torch.float32, device="cuda") + .reshape(6, 2) + .transpose(0, 1) + ) extra: dict[str, Any] = { MINIMAX_H3_LAYOUT_KEY: _layout(noncontiguous=True), "vsa_mode": "compete", @@ -82,7 +101,7 @@ def _encoded_batch(*, request_id: str | None = None) -> ForwardBatch: ) -def test_encoded_wire_state_is_minimal_cpu_contiguous_and_pickleable() -> None: +def test_encoded_wire_state_is_minimal_cuda_contiguous_and_pickleable() -> None: batch = _encoded_batch() assert not batch.prompt_embeds[0].is_contiguous() assert not batch.latents.is_contiguous() @@ -119,10 +138,10 @@ def test_encoded_wire_state_is_minimal_cpu_contiguous_and_pickleable() -> None: restored.layout.audio_indices, restored.layout.text_indices, ): - assert tensor.device.type == "cpu" + assert tensor.device.type == "cuda" assert tensor.is_contiguous() - roundtrip = restored.to_batch(device="cpu") + roundtrip = restored.to_batch() assert torch.equal(roundtrip.prompt_embeds[0], batch.prompt_embeds[0]) assert torch.equal(roundtrip.latents, batch.latents) assert torch.equal(roundtrip.audio_latents, batch.audio_latents) @@ -136,7 +155,7 @@ def test_encoded_wire_state_is_minimal_cpu_contiguous_and_pickleable() -> None: def test_denoised_wire_state_roundtrips_only_decode_inputs() -> None: encoded = MiniMaxH3EncodedState.from_batch(_encoded_batch(), request_id="request-8") - denoised_batch = encoded.to_batch(device="cpu") + denoised_batch = encoded.to_batch() denoised_batch.latents = denoised_batch.latents + 1 denoised_batch.audio_latents = denoised_batch.audio_latents - 1 @@ -168,9 +187,9 @@ def test_denoised_wire_state_roundtrips_only_decode_inputs() -> None: ({"raw_latent_shape": (1, 16, 0, 2, 2)}, "raw latent shape"), ({"num_inference_steps": 0}, "at least one denoising step"), ({"vsa_mode": "unknown"}, "vsa_mode"), - ({"video_latents": torch.zeros(3, 5)}, "video latent rows"), - ({"audio_latents": torch.zeros(1, 6)}, "audio latent rows"), - ({"prompt_embeds": torch.zeros(1, 2, 4)}, "prompt embedding rows"), + ({"video_latents": torch.zeros(3, 5, device="cuda")}, "video latent rows"), + ({"audio_latents": torch.zeros(1, 6, device="cuda")}, "audio latent rows"), + ({"prompt_embeds": torch.zeros(1, 2, 4, device="cuda")}, "prompt embedding rows"), ], ) def test_encoded_wire_state_rejects_incompatible_payloads(mutation: dict[str, Any], message: str) -> None: @@ -181,10 +200,10 @@ def test_encoded_wire_state_rejects_incompatible_payloads(mutation: dict[str, An def test_wire_state_rejects_noncontiguous_or_malformed_layout() -> None: state = MiniMaxH3EncodedState.from_batch(_encoded_batch(), request_id="valid") - with pytest.raises(ValueError, match="layout.position_ids must be a contiguous CPU tensor"): + with pytest.raises(ValueError, match="layout.position_ids must be a contiguous CUDA tensor"): replace(state, layout=_layout(noncontiguous=True)) - malformed = replace(_layout(), token_tags=torch.zeros(8, dtype=torch.long)) + malformed = replace(_layout(), token_tags=torch.zeros(8, dtype=torch.long, device="cuda")) with pytest.raises(ValueError, match=r"layout.token_tags.*expected \(9,\)"): replace(state, layout=malformed) diff --git a/fastvideo/worker/minimax_h3_disaggregated.py b/fastvideo/worker/minimax_h3_disaggregated.py index 0f536335f..fba6dafdb 100644 --- a/fastvideo/worker/minimax_h3_disaggregated.py +++ b/fastvideo/worker/minimax_h3_disaggregated.py @@ -33,6 +33,55 @@ from fastvideo.worker.ray_utils import assert_ray_available, ray logger = init_logger(__name__) +def _log_h3_state(name: str, state: Any) -> None: + total_bytes = 0 + + for field_name in ( + "prompt_embeds", + "video_latents", + "audio_latents", + ): + tensor = getattr(state, field_name, None) + if tensor is not None: + size = tensor.numel() * tensor.element_size() + total_bytes += size + logger.info( + "[RAY_PAYLOAD] %s.%s: %.2f MB shape=%s dtype=%s device=%s", + name, + field_name, + size / 1e6, + tuple(tensor.shape), + tensor.dtype, + tensor.device, + ) + + # layout contains tensors too. + layout = getattr(state, "layout", None) + if layout is not None: + for field_name in ( + "text_indices", + "video_indices", + "audio_indices", + ): + tensor = getattr(layout, field_name, None) + if tensor is not None: + size = tensor.numel() * tensor.element_size() + total_bytes += size + logger.info( + "[RAY_PAYLOAD] %s.layout.%s: %.2f MB shape=%s dtype=%s device=%s", + name, + field_name, + size / 1e6, + tuple(tensor.shape), + tensor.dtype, + tensor.device, + ) + + logger.info( + "[RAY_PAYLOAD] %s TOTAL: %.2f MB", + name, + total_bytes / 1e6, + ) def _node_resource(node_ip: str) -> str: return f"node:{node_ip}" @@ -146,22 +195,6 @@ class _MiniMaxH3DiTActor: self.pipeline = pipeline_cls(args.model_path, args) self.pipeline.post_init() - def denoise(self, state: MiniMaxH3EncodedState) -> MiniMaxH3DenoisedState: - return self.pipeline.denoise(state) - - def set_lora_adapter(self, lora_nickname: str, lora_path: str | None, strength: float, - accumulate: bool) -> dict[str, str]: - self.pipeline.set_lora_adapter(lora_nickname, lora_path, strength=strength, accumulate=accumulate) - return {"status": "lora_adapter_set"} - - def unmerge_lora_weights(self) -> dict[str, str]: - self.pipeline.unmerge_lora_weights() - return {"status": "lora_adapter_unmerged"} - - def merge_lora_weights(self) -> dict[str, str]: - self.pipeline.merge_lora_weights() - return {"status": "lora_adapter_merged"} - def health(self) -> dict[str, Any]: modules = tuple(sorted(self.pipeline.modules)) return { @@ -172,12 +205,11 @@ class _MiniMaxH3DiTActor: "all_resident": all(not is_lazy_module(module) for module in self.pipeline.modules.values()), } - def shutdown(self) -> dict[str, str]: - self.pipeline = None - cleanup_dist_env_and_memory(shutdown_ray=False) - return {"status": "shutdown_complete"} - - + def denoise(self, state: MiniMaxH3EncodedState) -> MiniMaxH3DenoisedState: + _log_h3_state("A_TO_B encoded", state) + result = self.pipeline.denoise(state) + _log_h3_state("B_TO_A denoised", result) + return result class RayMiniMaxH3DisaggregatedRuntime: """Own one persistent encoder/decoder actor and one persistent DiT actor."""