nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Remove CUDA to CPU and vice versa conversion

Browse Source 
And more network payload logging
This commit is contained in:
Satyam-53
2026-09-07 18:49:49 -07:00
committed by Satyam Srivastava
parent 4be9192615
commit 8d4e05acef
3 changed files with 122 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/pipelines/basic/minimax_h3/disaggregated.py
+31 -33
View File
@@ -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,
)
fastvideo/tests/worker/test_minimax_h3_disaggregated.py
+37 -18
View File
@@ -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)
fastvideo/worker/minimax_h3_disaggregated.py
+54 -22
View File
@@ -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."""