185 lines
7.1 KiB
Python
185 lines
7.1 KiB
Python
"""Run prompt-only H3 T2VA with Ulysses or TP+sequence parallelism."""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
|
from h3_blackwell_runtime.distributed import SequenceParallelContext
|
|
from h3_blackwell_runtime.packing import H3PromptPacker
|
|
from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
|
|
from h3_blackwell_runtime.sampler import sample_video_res_multistep
|
|
from h3_blackwell_runtime.t2v import random_av_latents
|
|
from h3_blackwell_runtime.tensor_parallel import configure_h3_tensor_parallel
|
|
from h3_blackwell_runtime.token_refiner import H3TokenRefiner
|
|
|
|
|
|
class ParallelDenoiser:
|
|
def __init__(self, model: H3PackedDenoiser, mode: str):
|
|
self.model = model
|
|
self.mode = mode
|
|
|
|
def __call__(self, hidden, timesteps, positions, segments, video_segment, audio_segment):
|
|
context = SequenceParallelContext.create(hidden.shape[0], heads=56, head_dim=128)
|
|
if self.mode == "ulysses":
|
|
return self.model.forward_sequence_parallel(
|
|
hidden, timesteps, positions, segments, video_segment, audio_segment, context,
|
|
)
|
|
return self.model.forward_tensor_parallel(
|
|
hidden, timesteps, positions, segments, video_segment, audio_segment, context,
|
|
)
|
|
|
|
|
|
def synchronize(device: torch.device) -> None:
|
|
torch.cuda.synchronize(device)
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--benchmark", type=Path, required=True)
|
|
parser.add_argument("--mode", choices=("ulysses", "tensor"), required=True)
|
|
parser.add_argument("--attention", default="sdpa")
|
|
parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
|
parser.add_argument("--text-encoder", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
|
|
parser.add_argument("--save-latent", type=Path)
|
|
parser.add_argument("--report", type=Path)
|
|
args = parser.parse_args()
|
|
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
torch.cuda.set_device(local_rank)
|
|
device = torch.device(f"cuda:{local_rank}")
|
|
dist.init_process_group("nccl", device_id=device)
|
|
rank = dist.get_rank()
|
|
world_size = dist.get_world_size()
|
|
benchmark = json.loads(args.benchmark.read_text(encoding="utf-8"))
|
|
|
|
load_started = time.perf_counter()
|
|
checkpoint = H3Checkpoint(args.model, device=device)
|
|
model = H3PackedDenoiser.from_checkpoint(
|
|
checkpoint, output_dtype=torch.bfloat16, attention_backend=args.attention,
|
|
).eval()
|
|
packer = H3PromptPacker(checkpoint)
|
|
if args.mode == "tensor":
|
|
partition_context = SequenceParallelContext.create(world_size, heads=56, head_dim=128)
|
|
configure_h3_tensor_parallel(model, partition_context)
|
|
synchronize(device)
|
|
model_load_seconds = time.perf_counter() - load_started
|
|
model_load_peak = torch.cuda.max_memory_allocated(device)
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
|
|
conditioning_started = time.perf_counter()
|
|
if rank == 0:
|
|
conditioner = Qwen3VLPromptConditioner(args.text_encoder, device=device, dtype=torch.float32)
|
|
refiner = H3TokenRefiner(checkpoint, attention_backend="sdpa").eval()
|
|
text = refiner(conditioner(benchmark["prompt"])).to(torch.bfloat16)
|
|
text_length = torch.tensor([text.shape[1]], dtype=torch.int64, device=device)
|
|
del conditioner, refiner
|
|
else:
|
|
text = None
|
|
text_length = torch.zeros(1, dtype=torch.int64, device=device)
|
|
checkpoint.release_cache()
|
|
torch.cuda.empty_cache()
|
|
dist.broadcast(text_length, src=0)
|
|
if rank != 0:
|
|
text = torch.empty(1, int(text_length.item()), 5376, dtype=torch.bfloat16, device=device)
|
|
dist.broadcast(text, src=0)
|
|
synchronize(device)
|
|
conditioning_seconds = time.perf_counter() - conditioning_started
|
|
conditioning_peak = torch.cuda.max_memory_allocated(device)
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
|
|
width, height = benchmark["resolution"]
|
|
video, audio, aligned_frames = random_av_latents(
|
|
width, height, benchmark["frames"], benchmark["seed"], device=device,
|
|
)
|
|
dist.barrier()
|
|
synchronize(device)
|
|
sampling_started = time.perf_counter()
|
|
video, audio = sample_video_res_multistep(
|
|
ParallelDenoiser(model, args.mode),
|
|
packer,
|
|
text,
|
|
video,
|
|
audio,
|
|
steps=benchmark["steps"],
|
|
seed=benchmark["seed"],
|
|
return_audio=True,
|
|
progress=rank == 0,
|
|
)
|
|
synchronize(device)
|
|
sampling_seconds = time.perf_counter() - sampling_started
|
|
|
|
timing = torch.tensor(
|
|
[model_load_seconds, conditioning_seconds, sampling_seconds],
|
|
dtype=torch.float64,
|
|
device=device,
|
|
)
|
|
timings = [torch.empty_like(timing) for _ in range(world_size)]
|
|
dist.all_gather(timings, timing)
|
|
checksums = torch.stack((video.float().sum(), audio.float().sum())).to(torch.float64)
|
|
all_checksums = [torch.empty_like(checksums) for _ in range(world_size)]
|
|
dist.all_gather(all_checksums, checksums)
|
|
checksum_stack = torch.stack(all_checksums)
|
|
if not torch.allclose(checksum_stack, checksum_stack[0].expand_as(checksum_stack), rtol=0, atol=1e-5):
|
|
raise RuntimeError(f"rank outputs diverged: {checksum_stack.cpu().tolist()}")
|
|
|
|
peak_memory = torch.tensor(
|
|
[model_load_peak, conditioning_peak, torch.cuda.max_memory_allocated(device)],
|
|
dtype=torch.int64,
|
|
device=device,
|
|
)
|
|
memory = [torch.empty_like(peak_memory) for _ in range(world_size)]
|
|
dist.all_gather(memory, peak_memory)
|
|
|
|
if rank == 0:
|
|
if args.save_latent is not None:
|
|
args.save_latent.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save({
|
|
"latent": video.cpu(),
|
|
"audio_latent": audio.cpu(),
|
|
"frames": aligned_frames,
|
|
"width": width,
|
|
"height": height,
|
|
"prompt": benchmark["prompt"],
|
|
"seed": benchmark["seed"],
|
|
"distributed_mode": args.mode,
|
|
"world_size": world_size,
|
|
"attention": args.attention,
|
|
}, args.save_latent)
|
|
timing_stack = torch.stack(timings).cpu()
|
|
report = {
|
|
"mode": args.mode,
|
|
"world_size": world_size,
|
|
"attention": args.attention,
|
|
"device": torch.cuda.get_device_name(device),
|
|
"torch": torch.__version__,
|
|
"benchmark": str(args.benchmark),
|
|
"resolution": [width, height],
|
|
"frames": aligned_frames,
|
|
"steps": benchmark["steps"],
|
|
"seed": benchmark["seed"],
|
|
"timings_max_rank_seconds": {
|
|
"model_load": float(timing_stack[:, 0].max()),
|
|
"conditioning": float(timing_stack[:, 1].max()),
|
|
"sampling": float(timing_stack[:, 2].max()),
|
|
},
|
|
"peak_allocated_bytes_by_rank": {
|
|
"model_load": [int(value[0].item()) for value in memory],
|
|
"conditioning": [int(value[1].item()) for value in memory],
|
|
"sampling": [int(value[2].item()) for value in memory],
|
|
},
|
|
"checksums": checksum_stack[0].cpu().tolist(),
|
|
"latent": str(args.save_latent) if args.save_latent is not None else None,
|
|
}
|
|
serialized = json.dumps(report, indent=2)
|
|
if args.report is not None:
|
|
args.report.parent.mkdir(parents=True, exist_ok=True)
|
|
args.report.write_text(serialized + "\n", encoding="utf-8")
|
|
print(serialized)
|
|
|
|
dist.destroy_process_group()
|