Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / encoder_parallel / transport.py: 96%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Transport CPU image shards without staging complete images on a device.""" 

16# pylint: disable=forbidden-backend-import 

17 

18from __future__ import annotations 

19 

20from typing import Any, Mapping 

21 

22import torch 

23import torch.distributed as dist 

24 

25from hyper_parallel.data.parallel import TPBatchBroadcaster 

26from hyper_parallel.distributed.encoder_parallel.contracts import EncoderLocalInputs 

27from hyper_parallel.distributed.encoder_parallel.planning import DecoderSequenceLayout, build_segments 

28from hyper_parallel.distributed.encoder_parallel.runtime import EncoderParallelPlan, EncoderParallelRuntime 

29 

30 

31_METADATA_FIELD = "_encoder_transport_metadata" 

32_FORWARD_FIELDS = ("encoder_parallel_plan", "encoder_parallel_inputs") 

33_CHUNK_BYTES = 4 * 1024 * 1024 

34 

35 

36class EncoderParallelBatchTransport: 

37 """Plan from small metadata and send only each destination's CPU patches. 

38 

39 Every decoder CP coordinate already has a complete CPU batch on TP0. TP 

40 peers may belong to different encoder DP buckets or CP shards, so their 

41 payloads are transported individually with bounded device staging. 

42 """ 

43 

44 def __init__(self, *, runtime: EncoderParallelRuntime, chunk_bytes: int = _CHUNK_BYTES) -> None: 

45 """Set the runtime and maximum device staging size in bytes.""" 

46 if isinstance(chunk_bytes, bool) or not isinstance(chunk_bytes, int) or chunk_bytes < 1: 

47 raise ValueError("encoder transport chunk_bytes must be a positive integer") 

48 self.runtime = runtime 

49 self.chunk_bytes = chunk_bytes 

50 

51 def _prepare_source(self, batch: Mapping[str, Any]) -> dict[str, Any]: 

52 """Validate CPU images and keep tensor payloads out of TP broadcasting.""" 

53 bridge = self.runtime.bridge 

54 prepared = {name: value for name, value in batch.items() if name not in bridge.input_fields} 

55 error = None 

56 items = () 

57 dtype = torch.float32 

58 try: 

59 if set(batch).intersection((_METADATA_FIELD, *_FORWARD_FIELDS)): 

60 raise ValueError("encoder transport fields cannot be supplied by the DataLoader") 

61 for name in bridge.input_fields: 

62 value = batch.get(name) 

63 if torch.is_tensor(value) and value.device.type != "cpu": 

64 raise ValueError(f"encoder source field {name} must remain on the CPU") 

65 items = bridge.parse_images(batch) 

66 pixels = batch.get(bridge.payload_field) 

67 if pixels is not None: 

68 dtype = pixels.dtype 

69 except (ValueError, TypeError, KeyError) as exc: 

70 error = str(exc) 

71 prepared[_METADATA_FIELD] = (items, dtype, error) 

72 return prepared 

73 

74 def _check_source_errors(self, error: str | None, device: Any) -> None: 

75 """Make all FSDP participants fail before plan or payload collectives.""" 

76 invalid = torch.tensor(int(error is not None), device=device) 

77 if dist.is_initialized() and self.runtime.global_mesh.size() > 1: 

78 dist.all_reduce(invalid, op=dist.ReduceOp.MAX, group=self.runtime.global_mesh.get_group()) 

79 if bool(invalid.item()): 

80 raise ValueError(f"invalid encoder CPU input: {error or 'another source rank rejected its batch'}") 

81 

82 def __call__(self, batch: Mapping[str, Any] | None, 

83 broadcaster: TPBatchBroadcaster) -> tuple[dict[str, Any], dict[str, Any]]: 

84 """Broadcast ordinary inputs, agree on a plan, and materialize local patches.""" 

85 source = self._prepare_source(batch) if batch is not None else None 

86 device_batch = broadcaster.broadcast(source, broadcast_all_fields=True) 

87 items, dtype, error = device_batch.pop(_METADATA_FIELD) 

88 self._check_source_errors(error, broadcaster.device) 

89 layout = DecoderSequenceLayout.build( 

90 mesh_context=self.runtime.mesh_context, cp_local_length=int(device_batch["input_ids"].shape[1]), 

91 ) 

92 plan = self.runtime.finalize_plan(items, layout) 

93 inputs = {"encoder_parallel_plan": plan} 

94 if plan.active: 

95 inputs["encoder_parallel_inputs"] = EncoderLocalInputs( 

96 self._transport_patches(batch, plan, dtype, broadcaster), 

97 ) 

98 return device_batch, inputs 

99 

100 def _transport_patches(self, batch: Mapping[str, Any] | None, plan: EncoderParallelPlan, 

101 dtype: torch.dtype, broadcaster: TPBatchBroadcaster) -> torch.Tensor: 

102 """Pack on the CPU and reuse a chunk-sized buffer across TP destinations.""" 

103 context = broadcaster.parallel_context 

104 local_count = sum(segment.local_padded_length for segment in plan.local_segments) 

105 local = torch.empty((local_count, *self.runtime.bridge.empty_payload_shape), 

106 dtype=dtype, device=broadcaster.device) 

107 chunk_elements = max(1, self.chunk_bytes // local.element_size()) 

108 target_segments = [] 

109 for target in range(context.tp_world_size): 

110 parent_rank = context.cp_rank * context.tp_world_size + target 

111 bucket, cp_rank = divmod(parent_rank, self.runtime.config.encoder_cp_size) 

112 target_segments.append(build_segments(plan.buckets[bucket], cp_rank, self.runtime.config.encoder_cp_size)) 

113 patch_elements = local[0].numel() 

114 max_elements = max((sum(segment.local_padded_length for segment in segments) * patch_elements 

115 for segments in target_segments[1:]), default=0) 

116 staging = torch.empty(min(chunk_elements, max_elements), dtype=dtype, device=broadcaster.device) 

117 for target, segments in enumerate(target_segments): 

118 packed = self.runtime.bridge.pack_payload(batch, segments) if context.tp_rank == 0 else None 

119 if target == 0: 

120 if context.tp_rank == 0: 

121 local.copy_(packed) 

122 continue 

123 count = sum(segment.local_padded_length for segment in segments) 

124 elements = count * patch_elements 

125 for start in range(0, elements, chunk_elements): 

126 length = min(chunk_elements, elements - start) 

127 buffer = (local.view(-1)[start:start + length] if context.tp_rank == target 

128 else staging[:length]) 

129 if context.tp_rank == 0: 

130 buffer.copy_(packed.view(-1)[start:start + length]) 

131 dist.broadcast(buffer, group=context.tp_group, group_src=0) 

132 return local