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()