263 lines
11 KiB
Python
263 lines
11 KiB
Python
|
import tempfile
|
||
|
|
from pathlib import Path
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.distributed as dist
|
||
|
|
import torch.multiprocessing as mp
|
||
|
|
from torch.nn import functional as F
|
||
|
|
|
||
|
|
from h3_blackwell_runtime.distributed import (
|
||
|
|
SequenceParallelContext,
|
||
|
|
balanced_ranges,
|
||
|
|
localize_segments,
|
||
|
|
range_lengths,
|
||
|
|
)
|
||
|
|
from h3_blackwell_runtime.final import H3FinalLayer
|
||
|
|
from h3_blackwell_runtime.nvfp4 import Nvfp4Linear, Nvfp4LinearTensors
|
||
|
|
from h3_blackwell_runtime.tensor_parallel import (
|
||
|
|
aligned_balanced_ranges,
|
||
|
|
select_nvfp4_outputs,
|
||
|
|
slice_nvfp4_inputs,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _init_gloo(rank: int, world_size: int, init_file: str) -> None:
|
||
|
|
dist.init_process_group(
|
||
|
|
"gloo",
|
||
|
|
init_method=f"file://{init_file}",
|
||
|
|
rank=rank,
|
||
|
|
world_size=world_size,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _transport_identity_worker(rank: int, world_size: int, init_file: str) -> None:
|
||
|
|
_init_gloo(rank, world_size, init_file)
|
||
|
|
try:
|
||
|
|
sequence, heads, head_dim = 17, 56, 3
|
||
|
|
context = SequenceParallelContext.create(sequence, heads, head_dim)
|
||
|
|
start, stop = context.local_token_range
|
||
|
|
full = torch.arange(sequence * heads * head_dim, dtype=torch.float32).reshape(1, sequence, heads, head_dim)
|
||
|
|
q = full[:, start:stop].contiguous()
|
||
|
|
k = q + 1_000_000
|
||
|
|
v = q + 2_000_000
|
||
|
|
full_q, full_k, full_v = context.seq_to_heads(q, k, v)
|
||
|
|
torch.testing.assert_close(context.heads_to_seq(full_q), q[0], rtol=0, atol=0)
|
||
|
|
torch.testing.assert_close(context.heads_to_seq(full_k), k[0], rtol=0, atol=0)
|
||
|
|
torch.testing.assert_close(context.heads_to_seq(full_v), v[0], rtol=0, atol=0)
|
||
|
|
finally:
|
||
|
|
dist.destroy_process_group()
|
||
|
|
|
||
|
|
|
||
|
|
def _attention_parity_worker(rank: int, world_size: int, init_file: str) -> None:
|
||
|
|
_init_gloo(rank, world_size, init_file)
|
||
|
|
try:
|
||
|
|
torch.manual_seed(440420)
|
||
|
|
sequence, heads, head_dim = 19, 56, 8
|
||
|
|
context = SequenceParallelContext.create(sequence, heads, head_dim)
|
||
|
|
q = torch.randn(1, sequence, heads, head_dim)
|
||
|
|
k = torch.randn_like(q)
|
||
|
|
v = torch.randn_like(q)
|
||
|
|
start, stop = context.local_token_range
|
||
|
|
local_q, local_k, local_v = context.seq_to_heads(
|
||
|
|
q[:, start:stop].contiguous(),
|
||
|
|
k[:, start:stop].contiguous(),
|
||
|
|
v[:, start:stop].contiguous(),
|
||
|
|
)
|
||
|
|
local_heads = F.scaled_dot_product_attention(
|
||
|
|
local_q.transpose(1, 2),
|
||
|
|
local_k.transpose(1, 2),
|
||
|
|
local_v.transpose(1, 2),
|
||
|
|
).transpose(1, 2)
|
||
|
|
actual = context.heads_to_seq(local_heads)
|
||
|
|
expected = F.scaled_dot_product_attention(
|
||
|
|
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
|
||
|
|
).transpose(1, 2)[0, start:stop]
|
||
|
|
torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
|
||
|
|
finally:
|
||
|
|
dist.destroy_process_group()
|
||
|
|
|
||
|
|
|
||
|
|
def _final_projection_worker(rank: int, world_size: int, init_file: str) -> None:
|
||
|
|
_init_gloo(rank, world_size, init_file)
|
||
|
|
try:
|
||
|
|
sequence, hidden_size = 9, 4
|
||
|
|
context = SequenceParallelContext.create(sequence, heads=4, head_dim=2)
|
||
|
|
hidden = torch.arange(sequence * hidden_size, dtype=torch.float32).reshape(sequence, hidden_size) / 10
|
||
|
|
timesteps = torch.tensor([0.25, 0.75])
|
||
|
|
layer = H3FinalLayer(
|
||
|
|
torch.zeros(1025, 2),
|
||
|
|
torch.ones(hidden_size),
|
||
|
|
torch.zeros(2 * hidden_size, 2),
|
||
|
|
torch.zeros(2 * hidden_size),
|
||
|
|
torch.arange(3 * hidden_size, dtype=torch.float32).reshape(3, hidden_size) / 10,
|
||
|
|
torch.tensor([0.1, 0.2, 0.3]),
|
||
|
|
torch.arange(2 * hidden_size, dtype=torch.float32).reshape(2, hidden_size) / 20,
|
||
|
|
torch.tensor([-0.1, 0.1]),
|
||
|
|
hidden_size=hidden_size,
|
||
|
|
)
|
||
|
|
video_segment = (3, 9, 0)
|
||
|
|
audio_segment = (0, 3, 1)
|
||
|
|
expected_video, expected_audio = layer(hidden, timesteps, video_segment, audio_segment)
|
||
|
|
start, stop = context.local_token_range
|
||
|
|
actual_video, actual_audio = layer.forward_sequence_parallel(
|
||
|
|
hidden[start:stop], timesteps, video_segment, audio_segment, context,
|
||
|
|
)
|
||
|
|
torch.testing.assert_close(actual_video, expected_video)
|
||
|
|
torch.testing.assert_close(actual_audio, expected_audio)
|
||
|
|
finally:
|
||
|
|
dist.destroy_process_group()
|
||
|
|
|
||
|
|
|
||
|
|
def _ragged_row_collectives_worker(rank: int, world_size: int, init_file: str) -> None:
|
||
|
|
_init_gloo(rank, world_size, init_file)
|
||
|
|
try:
|
||
|
|
sequence, features = 17, 5
|
||
|
|
context = SequenceParallelContext.create(sequence, heads=56, head_dim=2)
|
||
|
|
full = torch.arange(sequence * features, dtype=torch.float32).reshape(sequence, features)
|
||
|
|
start, stop = context.local_token_range
|
||
|
|
gathered = context.all_gather_rows(full[start:stop].contiguous())
|
||
|
|
torch.testing.assert_close(gathered, full, rtol=0, atol=0)
|
||
|
|
partial = full * float(rank + 1)
|
||
|
|
reduced = context.reduce_scatter_rows(partial)
|
||
|
|
expected = full[start:stop] * sum(range(1, world_size + 1))
|
||
|
|
torch.testing.assert_close(reduced, expected, rtol=0, atol=0)
|
||
|
|
finally:
|
||
|
|
dist.destroy_process_group()
|
||
|
|
|
||
|
|
|
||
|
|
def _tensor_parallel_math_worker(rank: int, world_size: int, init_file: str) -> None:
|
||
|
|
_init_gloo(rank, world_size, init_file)
|
||
|
|
try:
|
||
|
|
torch.manual_seed(440421)
|
||
|
|
sequence, hidden, heads, head_dim, intermediate = 17, 16, 14, 4, 224
|
||
|
|
context = SequenceParallelContext.create(sequence, heads=heads, head_dim=head_dim)
|
||
|
|
full_x = torch.randn(sequence, hidden)
|
||
|
|
start, stop = context.local_token_range
|
||
|
|
gathered_x = context.all_gather_rows(full_x[start:stop].contiguous())
|
||
|
|
|
||
|
|
inner = heads * head_dim
|
||
|
|
qkv_weight = torch.randn(3 * inner, hidden)
|
||
|
|
output_weight = torch.randn(hidden, inner)
|
||
|
|
output_bias = torch.randn(hidden)
|
||
|
|
head_start, head_stop = context.local_head_range
|
||
|
|
feature_start, feature_stop = head_start * head_dim, head_stop * head_dim
|
||
|
|
indices = torch.cat((
|
||
|
|
torch.arange(feature_start, feature_stop),
|
||
|
|
torch.arange(inner + feature_start, inner + feature_stop),
|
||
|
|
torch.arange(2 * inner + feature_start, 2 * inner + feature_stop),
|
||
|
|
))
|
||
|
|
local_qkv = F.linear(gathered_x, qkv_weight.index_select(0, indices))
|
||
|
|
local_inner = context.local_head_count * head_dim
|
||
|
|
local_q, local_k, local_v = local_qkv.split(local_inner, dim=-1)
|
||
|
|
local_q = local_q.view(1, sequence, context.local_head_count, head_dim)
|
||
|
|
local_k = local_k.view_as(local_q)
|
||
|
|
local_v = local_v.view_as(local_q)
|
||
|
|
local_attention = F.scaled_dot_product_attention(
|
||
|
|
local_q.transpose(1, 2), local_k.transpose(1, 2), local_v.transpose(1, 2),
|
||
|
|
).transpose(1, 2).reshape(sequence, local_inner)
|
||
|
|
partial_attention = F.linear(
|
||
|
|
local_attention, output_weight[:, feature_start:feature_stop],
|
||
|
|
)
|
||
|
|
actual_attention = context.reduce_scatter_rows(partial_attention) + output_bias
|
||
|
|
|
||
|
|
full_q, full_k, full_v = F.linear(full_x, qkv_weight).split(inner, dim=-1)
|
||
|
|
full_q = full_q.view(1, sequence, heads, head_dim)
|
||
|
|
full_k = full_k.view_as(full_q)
|
||
|
|
full_v = full_v.view_as(full_q)
|
||
|
|
expected_attention = F.linear(
|
||
|
|
F.scaled_dot_product_attention(
|
||
|
|
full_q.transpose(1, 2), full_k.transpose(1, 2), full_v.transpose(1, 2),
|
||
|
|
).transpose(1, 2).reshape(sequence, inner),
|
||
|
|
output_weight,
|
||
|
|
output_bias,
|
||
|
|
)[start:stop]
|
||
|
|
torch.testing.assert_close(actual_attention, expected_attention, rtol=2e-5, atol=2e-5)
|
||
|
|
|
||
|
|
fc1_weight = torch.randn(2 * intermediate, hidden)
|
||
|
|
fc2_weight = torch.randn(hidden, intermediate)
|
||
|
|
fc2_bias = torch.randn(hidden)
|
||
|
|
mlp_ranges = aligned_balanced_ranges(intermediate, world_size, 32)
|
||
|
|
mlp_start, mlp_stop = mlp_ranges[rank]
|
||
|
|
local_fc1_weight = torch.cat((
|
||
|
|
fc1_weight[mlp_start:mlp_stop],
|
||
|
|
fc1_weight[intermediate + mlp_start:intermediate + mlp_stop],
|
||
|
|
))
|
||
|
|
gate, up = F.linear(gathered_x, local_fc1_weight).chunk(2, dim=-1)
|
||
|
|
partial_mlp = F.linear(F.silu(gate) * up, fc2_weight[:, mlp_start:mlp_stop])
|
||
|
|
actual_mlp = context.reduce_scatter_rows(partial_mlp) + fc2_bias
|
||
|
|
full_gate, full_up = F.linear(full_x, fc1_weight).chunk(2, dim=-1)
|
||
|
|
expected_mlp = F.linear(F.silu(full_gate) * full_up, fc2_weight, fc2_bias)[start:stop]
|
||
|
|
# TP reduction changes FP32 accumulation order across rank partials.
|
||
|
|
torch.testing.assert_close(actual_mlp, expected_mlp, rtol=1e-4, atol=2e-4)
|
||
|
|
finally:
|
||
|
|
dist.destroy_process_group()
|
||
|
|
|
||
|
|
|
||
|
|
def _run_distributed(worker, world_size: int) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as directory:
|
||
|
|
init_file = str(Path(directory) / "process-group")
|
||
|
|
mp.spawn(worker, args=(world_size, init_file), nprocs=world_size, join=True)
|
||
|
|
|
||
|
|
|
||
|
|
class DistributedPartitionContracts(unittest.TestCase):
|
||
|
|
def test_balanced_ragged_ranges(self):
|
||
|
|
ranges = balanced_ranges(56, 6)
|
||
|
|
self.assertEqual(range_lengths(ranges), (10, 10, 9, 9, 9, 9))
|
||
|
|
self.assertEqual(ranges[0], (0, 10))
|
||
|
|
self.assertEqual(ranges[-1], (47, 56))
|
||
|
|
|
||
|
|
def test_ranges_reject_empty_partitions(self):
|
||
|
|
with self.assertRaisesRegex(ValueError, "non-empty"):
|
||
|
|
balanced_ranges(3, 4)
|
||
|
|
|
||
|
|
def test_segments_are_clipped_and_rebased(self):
|
||
|
|
segments = [(0, 4, 1), (4, 10, 2), (10, 15, 3)]
|
||
|
|
self.assertEqual(localize_segments(segments, 3, 12), [(0, 1, 1), (1, 7, 2), (7, 9, 3)])
|
||
|
|
|
||
|
|
def test_transport_identity_for_planned_world_sizes(self):
|
||
|
|
for world_size in (2, 4, 6, 8):
|
||
|
|
with self.subTest(world_size=world_size):
|
||
|
|
_run_distributed(_transport_identity_worker, world_size)
|
||
|
|
|
||
|
|
def test_two_rank_sdpa_matches_single_process(self):
|
||
|
|
_run_distributed(_attention_parity_worker, 2)
|
||
|
|
|
||
|
|
def test_distributed_final_projection_matches_single_process(self):
|
||
|
|
_run_distributed(_final_projection_worker, 2)
|
||
|
|
|
||
|
|
def test_ragged_all_gather_and_reduce_scatter(self):
|
||
|
|
_run_distributed(_ragged_row_collectives_worker, 6)
|
||
|
|
|
||
|
|
def test_tensor_parallel_attention_and_mlp_match_dense_math(self):
|
||
|
|
for world_size in (2, 6):
|
||
|
|
with self.subTest(world_size=world_size):
|
||
|
|
_run_distributed(_tensor_parallel_math_worker, world_size)
|
||
|
|
|
||
|
|
def test_nvfp4_column_and_row_shards_preserve_layout(self):
|
||
|
|
tensors = Nvfp4LinearTensors(
|
||
|
|
weight=torch.arange(96 * 32, dtype=torch.int32).to(torch.uint8).reshape(96, 32),
|
||
|
|
weight_scale=torch.arange(96 * 4, dtype=torch.float32).to(torch.float8_e4m3fn).reshape(96, 4),
|
||
|
|
weight_scale_2=torch.tensor(0.5),
|
||
|
|
bias=torch.arange(96, dtype=torch.bfloat16),
|
||
|
|
pre_quant_scale=torch.arange(64, dtype=torch.bfloat16),
|
||
|
|
full_precision_matrix_mult=False,
|
||
|
|
in_features=64,
|
||
|
|
out_features=96,
|
||
|
|
)
|
||
|
|
linear = Nvfp4Linear(tensors)
|
||
|
|
column = select_nvfp4_outputs(linear, ((0, 16), (32, 48)))
|
||
|
|
self.assertEqual(tuple(column.weight.shape), (32, 32))
|
||
|
|
self.assertEqual(tuple(column.weight_scale.shape), (32, 4))
|
||
|
|
self.assertEqual(column.out_features, 32)
|
||
|
|
row, bias = slice_nvfp4_inputs(linear, 32, 64)
|
||
|
|
self.assertEqual(tuple(row.weight.shape), (96, 16))
|
||
|
|
self.assertEqual(tuple(row.weight_scale.shape), (96, 2))
|
||
|
|
self.assertEqual(row.in_features, 32)
|
||
|
|
self.assertIsNone(row.bias)
|
||
|
|
torch.testing.assert_close(bias, tensors.bias)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|