105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
"""Packed-input H3 transformer core, independent of ComfyUI node execution."""
|
|
|
|
import torch
|
|
from torch import nn
|
|
from typing import TYPE_CHECKING
|
|
|
|
from .attention import DEFAULT_ATTENTION_BACKEND
|
|
from .backbone import H3DenoiserBackbone
|
|
from .checkpoint import H3Checkpoint
|
|
from .final import H3FinalLayer
|
|
|
|
if TYPE_CHECKING:
|
|
from .distributed import SequenceParallelContext
|
|
|
|
|
|
class H3PackedDenoiser(nn.Module):
|
|
"""Run the H3 transformer once its Ref2VA payload has been packed into hidden rows."""
|
|
|
|
def __init__(self, backbone: H3DenoiserBackbone, final_layer: H3FinalLayer):
|
|
super().__init__()
|
|
self.backbone = backbone
|
|
self.final_layer = final_layer
|
|
|
|
@classmethod
|
|
def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = DEFAULT_ATTENTION_BACKEND):
|
|
return cls(
|
|
H3DenoiserBackbone.from_checkpoint(checkpoint, output_dtype=output_dtype, attention_backend=attention_backend),
|
|
H3FinalLayer.from_checkpoint(checkpoint, output_dtype=output_dtype),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden: torch.Tensor,
|
|
timesteps: torch.Tensor,
|
|
position_ids: torch.Tensor,
|
|
segments: list[tuple[int, int, int]],
|
|
video_segment: tuple[int, int, int],
|
|
audio_segment: tuple[int, int, int],
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
hidden = self.backbone(hidden, timesteps, position_ids, segments)
|
|
return self.final_layer(hidden, timesteps, video_segment, audio_segment)
|
|
|
|
def forward_sequence_parallel(
|
|
self,
|
|
full_hidden: torch.Tensor,
|
|
timesteps: torch.Tensor,
|
|
full_position_ids: torch.Tensor,
|
|
segments: list[tuple[int, int, int]],
|
|
video_segment: tuple[int, int, int],
|
|
audio_segment: tuple[int, int, int],
|
|
context: "SequenceParallelContext",
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Run all 50 blocks with token-sharded activations and Ulysses attention."""
|
|
if full_hidden.shape[0] != context.sequence_length:
|
|
raise ValueError("packed hidden length does not match sequence-parallel context")
|
|
start, stop = context.local_token_range
|
|
local_hidden = full_hidden[start:stop].contiguous()
|
|
local_positions = full_position_ids[start:stop].contiguous()
|
|
del full_hidden, full_position_ids
|
|
local_hidden = self.backbone(
|
|
local_hidden,
|
|
timesteps,
|
|
local_positions,
|
|
segments,
|
|
sequence_parallel=context,
|
|
)
|
|
return self.final_layer.forward_sequence_parallel(
|
|
local_hidden,
|
|
timesteps,
|
|
video_segment,
|
|
audio_segment,
|
|
context,
|
|
)
|
|
|
|
def forward_tensor_parallel(
|
|
self,
|
|
full_hidden: torch.Tensor,
|
|
timesteps: torch.Tensor,
|
|
full_position_ids: torch.Tensor,
|
|
segments: list[tuple[int, int, int]],
|
|
video_segment: tuple[int, int, int],
|
|
audio_segment: tuple[int, int, int],
|
|
context: "SequenceParallelContext",
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Run TP-sharded NVFP4 linears with ragged sequence-sharded residuals."""
|
|
if full_hidden.shape[0] != context.sequence_length:
|
|
raise ValueError("packed hidden length does not match tensor-parallel context")
|
|
start, stop = context.local_token_range
|
|
local_hidden = full_hidden[start:stop].contiguous()
|
|
local_positions = full_position_ids[start:stop].contiguous()
|
|
del full_hidden, full_position_ids
|
|
local_hidden = self.backbone(
|
|
local_hidden,
|
|
timesteps,
|
|
local_positions,
|
|
segments,
|
|
tensor_parallel=context,
|
|
)
|
|
return self.final_layer.forward_sequence_parallel(
|
|
local_hidden,
|
|
timesteps,
|
|
video_segment,
|
|
audio_segment,
|
|
context,
|
|
)
|