From 47b542b951698490a9ac60e8fabe5450aea1395d Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 11 Mar 2026 16:43:55 +0100 Subject: [PATCH 1/5] AI draft i'll go from that --- src/transformers/core_model_loading.py | 30 +- src/transformers/integrations/__init__.py | 14 + .../integrations/pipeline_parallel.py | 427 ++++++++++++++++++ tests/pipeline_parallel_smoke.py | 67 +++ 4 files changed, 536 insertions(+), 2 deletions(-) create mode 100644 src/transformers/integrations/pipeline_parallel.py create mode 100644 tests/pipeline_parallel_smoke.py diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 32dce7c25f70..9b3b7031b3f6 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -32,6 +32,7 @@ import torch from .integrations.accelerate import get_device, offload_weight +from .integrations.pipeline_parallel import flatten_pp_units, partition_units from .integrations.tensor_parallel import ALL_PARALLEL_STYLES from .utils import is_env_variable_true from .utils.loading_report import LoadStateDictInfo @@ -1116,12 +1117,29 @@ def convert_and_load_state_dict_in_model( else: thread_pool = ThreadPoolExecutor(max_workers=GLOBAL_WORKERS) + # Pipeline parallel partitioning (if available) + pp_plan = model.pp_plan if getattr(model, "supports_pp_plan", False) else {} + local_pp_units: set[str] = set() + if device_mesh is not None and pp_plan: + pp_mesh = ( + device_mesh["pp"] + if device_mesh.ndim > 1 and "pp" in (device_mesh.mesh_dim_names or {}) + else device_mesh + ) + pp_size = pp_mesh.size() + pp_rank = pp_mesh.get_local_rank() + pp_units = flatten_pp_units(model, pp_plan) + pp_partitions = partition_units(pp_units, pp_size) + local_pp_units = set(pp_partitions[pp_rank]) if pp_rank < len(pp_partitions) else set() + renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)] converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)] param_name_to_load: dict[str, WeightRenaming | WeightConverter] = {} # build '(?P.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'} # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched. + tp_plan_alt = None + tp_plan_by_group_name = {} if tp_plan != {}: tp_plan_alt, tp_plan_by_group_name, _ = build_glob_alternation(list(tp_plan.keys())) if dtype_plan != {}: @@ -1179,7 +1197,12 @@ def convert_and_load_state_dict_in_model( # 4. Handle TP sharding or device_map placement future_or_tensor = None - if device_mesh: + param_module = renamed_key.rsplit(".", 1)[0] + is_local_pp = True if len(local_pp_units) == 0 else any( + param_module.startswith(unit) for unit in local_pp_units + ) + + if device_mesh and is_local_pp and tp_plan_alt is not None: if matched_tp_pattern := tp_plan_alt.search(renamed_key): matched_tp_pattern = tp_plan_by_group_name[matched_tp_pattern.lastgroup] if getattr(mapping, "distributed_operation", None) is None: @@ -1202,7 +1225,10 @@ def convert_and_load_state_dict_in_model( ) if future_or_tensor is None: - param_device = get_device(device_map, renamed_key, valid_torch_device=True) + if is_local_pp: + param_device = get_device(device_map, renamed_key, valid_torch_device=True) + else: + param_device = "cpu" future_or_tensor = spawn_materialize(thread_pool, tensor, param_device, _dtype) mapping.add_tensor(renamed_key, original_key, source_pattern, future_or_tensor) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 246790a13baf..3b2d0ef0d252 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -163,6 +163,13 @@ "ALL_PARALLEL_STYLES", "translate_to_torch_parallel_style", ] +_import_structure["pipeline_parallel"] = [ + "PipelineParallelSkipped", + "initialize_pipeline_parallelism", + "add_pipeline_parallel_hooks", + "flatten_pp_units", + "partition_units", +] try: if not is_torch_greater_or_equal("2.5"): raise OptionalDependencyNotAvailable() @@ -306,6 +313,13 @@ shard_and_distribute_module, translate_to_torch_parallel_style, ) + from .pipeline_parallel import ( + PipelineParallelSkipped, + add_pipeline_parallel_hooks, + flatten_pp_units, + initialize_pipeline_parallelism, + partition_units, + ) try: if not is_torch_greater_or_equal("2.5"): diff --git a/src/transformers/integrations/pipeline_parallel.py b/src/transformers/integrations/pipeline_parallel.py new file mode 100644 index 000000000000..8a5d4455a29b --- /dev/null +++ b/src/transformers/integrations/pipeline_parallel.py @@ -0,0 +1,427 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import math +import pickle +from contextlib import contextmanager +from typing import Iterable, Sequence + +from ..utils import logging +from ..utils.import_utils import is_torch_available + + +if is_torch_available(): + import torch + import torch.distributed as dist + from torch import nn + from torch.autograd.function import Function + + +logger = logging.get_logger(__name__) + + +class PipelineParallelSkipped(nn.Module): + """Placeholder module used on ranks that do not own a stage. + + It simply returns the first positional argument or the value of + ``hidden_states`` if provided. This keeps the module graph valid while + bypassing the computation on non-local pipeline stages. + """ + + def __init__(self): + super().__init__() + + def forward(self, *args, **kwargs): + if "hidden_states" in kwargs: + return kwargs["hidden_states"] + return args[0] if len(args) > 0 else None + + +@contextmanager +def _silent_non_primary(local_rank: int): + """Silence stdout/stderr on non-primary ranks to reduce clutter.""" + + if local_rank == 0: + yield + else: + import os + import sys + + old_stdout, old_stderr = sys.stdout, sys.stderr + sys.stdout = open(os.devnull, "w") + sys.stderr = open(os.devnull, "w") + try: + yield + finally: + sys.stdout.close() + sys.stderr.close() + sys.stdout, sys.stderr = old_stdout, old_stderr + + +def initialize_pipeline_parallelism(pp_plan: dict[str, tuple[list[str], list[str]]] | None, pp_size: int | None = None, device_mesh=None, device_map=None): + """Infer pipeline parallel rank/size and device mapping. + + This mirrors ``initialize_tensor_parallelism`` but only sets PP-related + attributes. When a multi-dimensional ``device_mesh`` is provided, the + ``"pp"`` dimension is used; otherwise ``WORLD_SIZE``/``RANK`` env vars are + consulted. + """ + + if not is_torch_available(): + raise ImportError("PyTorch is required for pipeline parallelism") + + if device_mesh is not None: + # If the mesh is multi-dim, use the "pp" named dimension when present. + if device_mesh.ndim > 1 and "pp" in (device_mesh.mesh_dim_names or {}): + pp_mesh = device_mesh["pp"] + else: + pp_mesh = device_mesh + pp_size = pp_mesh.size() + pp_rank = pp_mesh.get_local_rank() + return device_map, pp_mesh, pp_size, pp_rank + + # Fallback to env-derived initialization (single-dimension pipeline only) + if not dist.is_available(): + raise RuntimeError("torch.distributed is required for pipeline parallelism") + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before using pipeline parallelism") + + world_size = dist.get_world_size() + rank = dist.get_rank() + pp_size = pp_size or world_size + if world_size != pp_size: + raise ValueError("When no device_mesh is provided, world_size must equal pp_size") + return device_map, None, pp_size, rank + + +def _expand_module_name(model: nn.Module, name: str) -> list[str]: + """Expand ModuleList/Sequential entries into fully qualified names.""" + + try: + module = model.get_submodule(name) + except Exception: + return [name] + + if isinstance(module, (nn.ModuleList, nn.Sequential)): + return [f"{name}.{idx}" for idx, _ in enumerate(module)] + return [name] + + +def flatten_pp_units(model: nn.Module, pp_plan: dict[str, tuple[list[str], list[str]]]) -> list[tuple[str, nn.Module, list[str], list[str]]]: + """Return ordered pipeline units expanded from ``pp_plan`` keys.""" + + units = [] + for key, (inputs, outputs) in pp_plan.items(): + for expanded in _expand_module_name(model, key): + try: + mod = model.get_submodule(expanded) + except Exception: + continue + units.append((expanded, mod, inputs, outputs)) + return units + + +def partition_units(units: list[tuple[str, nn.Module, list[str], list[str]]], pp_size: int) -> list[list[str]]: + """Evenly partition units into ``pp_size`` contiguous chunks.""" + + if pp_size <= 0: + raise ValueError("pp_size must be positive") + chunk = math.ceil(len(units) / pp_size) if len(units) > 0 else 1 + partitions = [] + for idx in range(pp_size): + start = idx * chunk + stop = min((idx + 1) * chunk, len(units)) + partitions.append([name for name, _, _, _ in units[start:stop]]) + return partitions + + +def _group_ranks(group): + """Return (local_rank, global_rank, world_size) for a process group or the default group.""" + + local_rank = dist.get_rank(group) if group is not None else dist.get_rank() + global_rank = dist.get_global_rank(group, local_rank) if group is not None else local_rank + world_size = dist.get_world_size(group) if group is not None else dist.get_world_size() + return local_rank, global_rank, world_size + + +def _use_cuda_stream(tensors: Sequence[torch.Tensor]) -> bool: + return any(t.is_cuda for t in tensors) and torch.cuda.is_available() + + +def _global_to_group_rank(group, global_rank: int) -> int: + return dist.get_group_rank(group, global_rank) if group is not None else global_rank + + +def send_tensors(tensors: Sequence[torch.Tensor], dst: int, group=None, use_cuda_stream: bool = True) -> None: + """Send tensors with shape/dtype metadata to ``dst``.""" + + if len(tensors) == 0: + return + meta = [(tuple(t.shape), t.dtype) for t in tensors] + meta_bytes = pickle.dumps(meta) + meta_tensor = torch.tensor(list(meta_bytes), dtype=torch.uint8) + meta_len = torch.tensor([meta_tensor.numel()], dtype=torch.int64) + dist.send(meta_len, dst=dst, group=group) + dist.send(meta_tensor, dst=dst, group=group) + + stream = torch.cuda.Stream() if use_cuda_stream and _use_cuda_stream(tensors) else None + handles = [] + if stream is not None: + with torch.cuda.stream(stream): + for tensor in tensors: + handles.append(dist.isend(tensor, dst=dst, group=group)) + torch.cuda.current_stream().wait_stream(stream) + else: + for tensor in tensors: + handles.append(dist.isend(tensor, dst=dst, group=group)) + for h in handles: + h.wait() + + +def recv_tensors(count: int, src: int, device, group=None, use_cuda_stream: bool = True) -> list[torch.Tensor]: + """Receive ``count`` tensors from ``src`` using metadata first.""" + + meta_len = torch.empty(1, dtype=torch.int64) + dist.recv(meta_len, src=src, group=group) + meta_tensor = torch.empty(int(meta_len.item()), dtype=torch.uint8) + dist.recv(meta_tensor, src=src, group=group) + meta = pickle.loads(bytes(meta_tensor.tolist())) + tensors: list[torch.Tensor] = [] + stream = torch.cuda.Stream() if use_cuda_stream and torch.cuda.is_available() and device is not None and torch.device(device).type == "cuda" else None + handles = [] + for shape, dtype in meta: + buffer = torch.empty(shape, device=device, dtype=dtype) + tensors.append(buffer) + if stream is not None: + with torch.cuda.stream(stream): + handles.append(dist.irecv(buffer, src=src, group=group)) + else: + handles.append(dist.irecv(buffer, src=src, group=group)) + if stream is not None: + torch.cuda.current_stream().wait_stream(stream) + for h in handles: + h.wait() + return tensors + + +def send_grads(tensors: Sequence[torch.Tensor], dst: int, group=None, use_cuda_stream: bool = True) -> None: + """Send gradients to previous stage.""" + + filtered = [t for t in tensors if t is not None] + if len(filtered) == 0: + return + send_tensors(filtered, dst, group=group, use_cuda_stream=use_cuda_stream) + + +def recv_grads(count: int, src: int, device, group=None, use_cuda_stream: bool = True) -> list[torch.Tensor]: + """Receive gradients from next stage.""" + + return recv_tensors(count, src, device, group=group, use_cuda_stream=use_cuda_stream) + + +class _PipelineRecvFn(Function): + @staticmethod + def forward(ctx, count: int, src: int, device, group, use_cuda_stream: bool): + ctx.src = src + ctx.group = group + ctx.use_cuda_stream = use_cuda_stream + ctx.device = device + tensors = recv_tensors(count, src=src, device=device, group=group, use_cuda_stream=use_cuda_stream) + if count == 1: + return tensors[0] + return tuple(tensors) + + @staticmethod + def backward(ctx, *grad_outputs): + grads = [g for g in grad_outputs] + send_grads(grads, dst=ctx.src, group=ctx.group, use_cuda_stream=ctx.use_cuda_stream) + # None gradients for non-tensor args + return (None, None, None, None, None) + + +class _PipelineSendFn(Function): + @staticmethod + def forward(ctx, tensor: torch.Tensor, dst: int, group, use_cuda_stream: bool): + ctx.dst = dst + ctx.group = group + ctx.use_cuda_stream = use_cuda_stream + send_tensors([tensor], dst=dst, group=group, use_cuda_stream=use_cuda_stream) + return tensor + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grads = recv_grads(1, src=ctx.dst, device=grad_output.device, group=ctx.group, use_cuda_stream=ctx.use_cuda_stream) + return grads[0], None, None, None + + +def _pipeline_coordinates(device_mesh): + """Return (pp_rank, pp_size, tp_rank) tuple for the current process.""" + + if device_mesh is None: + return dist.get_rank(), dist.get_world_size(), 0 + mesh_names = device_mesh.mesh_dim_names or {} + coord = device_mesh.get_coordinate() + if coord is None: + raise RuntimeError("Device mesh coordinate is undefined") + if "pp" in mesh_names: + pp_dim = mesh_names.index("pp") if isinstance(mesh_names, tuple) else list(mesh_names).index("pp") + else: + pp_dim = 0 + if device_mesh.ndim > 1 and "tp" in mesh_names: + tp_dim = list(mesh_names).index("tp") if not isinstance(mesh_names, dict) else mesh_names.index("tp") + tp_rank = coord[tp_dim] + else: + tp_rank = 0 + return coord[pp_dim], device_mesh.size(pp_dim) if device_mesh is not None else dist.get_world_size(), tp_rank + + +def add_pipeline_parallel_hooks( + model: nn.Module, + pp_plan: dict[str, tuple[list[str], list[str]]], + device_mesh=None, + pp_size: int | None = None, + use_cuda_stream: bool = True, +): + """Attach send/recv hooks according to ``pp_plan``. + + This mirrors tensor-parallel's hook-based API: modules are prepared once, + and normal ``forward`` calls will automatically exchange activations + between pipeline stages. + """ + + if pp_plan is None or len(pp_plan) == 0: + return model + if not is_torch_available(): + raise ImportError("PyTorch is required for pipeline parallelism") + if not dist.is_available() or not dist.is_initialized(): + logger.warning("torch.distributed is not initialized; running without pipeline parallelism") + return model + + _, pp_mesh, pp_size, pp_rank = initialize_pipeline_parallelism(pp_plan, pp_size, device_mesh=device_mesh) + pp_group = pp_mesh.get_group() if pp_mesh is not None else None + model._pp_size = pp_size + model._pp_rank = pp_rank + model._pp_group = pp_group + + units = flatten_pp_units(model, pp_plan) + partitions = partition_units(units, pp_size) + local_unit_names = set(partitions[pp_rank]) if pp_rank < len(partitions) else set() + + # Replace non-local modules with skipped placeholders to avoid computation + for name, module, _, _ in units: + if name not in local_unit_names: + parent_path, child_name = name.rsplit(".", 1) if "." in name else ("", name) + parent = model.get_submodule(parent_path) if parent_path != "" else model + setattr(parent, child_name, PipelineParallelSkipped()) + + # Identify first/last local modules to add comm hooks + local_units_ordered = [u for u in units if u[0] in local_unit_names] + if len(local_units_ordered) == 0: + return model + first_name, first_mod, first_inputs, _ = local_units_ordered[0] + last_name, last_mod, _, last_outputs = local_units_ordered[-1] + + runtime_key = "_hf_pp_runtime" + + def _init_runtime(module, args, kwargs): + # args/kwargs are from model forward; seed activation dict with kwargs + runtime = {"activations": {}, "device": args[0].device if len(args) > 0 and hasattr(args[0], "device") else None} + runtime["activations"].update(kwargs) + setattr(model, runtime_key, runtime) + return None + + def _cleanup_runtime(module, args, kwargs, output): + if hasattr(model, runtime_key): + delattr(model, runtime_key) + return output + + model.register_forward_pre_hook(_init_runtime, with_kwargs=True) + model.register_forward_hook(_cleanup_runtime, with_kwargs=True) + + def _recv_pre_hook(mod, args, kwargs): + runtime = getattr(model, runtime_key) + local_pp_rank, global_pp_rank, _ = _group_ranks(pp_group) + if pp_rank > 0: + src_local = local_pp_rank - 1 + src_global = dist.get_global_rank(pp_group, src_local) if pp_group is not None else global_pp_rank - 1 + with _silent_non_primary(pp_rank): + tensors = [ + _PipelineRecvFn.apply( + 1, + src_global, + runtime.get("device"), + pp_group, + use_cuda_stream, + ) + for _ in first_inputs + ] + for name, tensor in zip(first_inputs, tensors): + runtime["activations"][name] = tensor + return tuple(tensors), {} + else: + return args, kwargs + # Build kwargs for this module from runtime activations + new_kwargs = {name: runtime["activations"].get(name) for name in first_inputs if name in runtime["activations"]} + return (), new_kwargs + + def _send_post_hook(mod, args, kwargs, output): + runtime = getattr(model, runtime_key) + outputs = output if isinstance(output, (tuple, list)) else (output,) + to_send = [] + for name in last_outputs: + # Map sequentially to module outputs + idx = last_outputs.index(name) + if idx < len(outputs): + tensor = outputs[idx] + else: + tensor = outputs[0] + runtime["activations"][name] = tensor + to_send.append(tensor) + + if pp_rank < pp_size - 1: + local_pp_rank, global_pp_rank, _ = _group_ranks(pp_group) + dst_local = local_pp_rank + 1 + dst_global = dist.get_global_rank(pp_group, dst_local) if pp_group is not None else global_pp_rank + 1 + to_send = [_PipelineSendFn.apply(tensor, dst_global, pp_group, use_cuda_stream) for tensor in to_send] + return output + + first_mod.register_forward_pre_hook(_recv_pre_hook, with_kwargs=True) + last_mod.register_forward_hook(_send_post_hook, with_kwargs=True) + + # For all local modules, store outputs into activation dict for downstream + def _store_outputs(mod, args, kwargs, output, output_names): + runtime = getattr(model, runtime_key) + outputs = output if isinstance(output, (tuple, list)) else (output,) + for idx, name in enumerate(output_names): + if idx < len(outputs): + runtime["activations"][name] = outputs[idx] + # For first module when not pipeline root, ensure gradients propagate + if pp_rank > 0 and outputs[idx].requires_grad and name in first_inputs: + outputs[idx].register_hook(lambda grad, name=name: grad) + return output + + for name, module, _, out_names in local_units_ordered: + module.register_forward_hook(lambda m, a, kw, o, out_names=out_names: _store_outputs(m, a, kw, o, out_names), with_kwargs=True) + + return model + + +__all__ = [ + "PipelineParallelSkipped", + "initialize_pipeline_parallelism", + "add_pipeline_parallel_hooks", + "flatten_pp_units", + "partition_units", +] diff --git a/tests/pipeline_parallel_smoke.py b/tests/pipeline_parallel_smoke.py new file mode 100644 index 000000000000..9c79641ffe86 --- /dev/null +++ b/tests/pipeline_parallel_smoke.py @@ -0,0 +1,67 @@ +import os +import sys + +import torch +import torch.distributed as dist +import torch.nn as nn + +from transformers.integrations.pipeline_parallel import add_pipeline_parallel_hooks + + +class TinyModel(nn.Module): + def __init__(self, hidden=8): + super().__init__() + self.lin0 = nn.Linear(hidden, hidden, bias=False) + self.lin1 = nn.Linear(hidden, hidden, bias=False) + + def forward(self, x): + x = self.lin0(x) + x = torch.relu(x) + x = self.lin1(x) + return x + + +def init_dist(): + backend = "nccl" if torch.cuda.is_available() else "gloo" + dist.init_process_group(backend=backend) + if torch.cuda.is_available(): + torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count()) + + +def main(): + init_dist() + rank = dist.get_rank() + device = torch.device("cuda", rank % torch.cuda.device_count()) if torch.cuda.is_available() else torch.device("cpu") + + torch.manual_seed(0) + base_model = TinyModel().to("cpu") + state_dict = base_model.state_dict() + dist.broadcast_object_list([state_dict], src=0) + + model = TinyModel().to(device) + model.load_state_dict(state_dict) + + pp_plan = { + "lin0": (["x"], ["x"]), + "lin1": (["x"], ["x"]), + } + add_pipeline_parallel_hooks(model, pp_plan, device_mesh=None, pp_size=dist.get_world_size()) + + x = torch.randn(2, 8, device=device) + out = model(x) + + if rank == dist.get_world_size() - 1: + ref = base_model(x.to("cpu")).detach() + dist.send(out.detach().cpu(), dst=0) + if rank == 0 and dist.get_world_size() == 1: # single rank debug + torch.testing.assert_close(out.cpu(), ref, atol=1e-4, rtol=1e-4) + if rank == 0 and dist.get_world_size() > 1: + recv = torch.empty_like(x.cpu()) + dist.recv(recv, src=dist.get_world_size() - 1) + ref = base_model(x.to("cpu")).detach() + torch.testing.assert_close(recv, ref, atol=1e-4, rtol=1e-4) + print("pipeline_parallel_smoke: OK") + + +if __name__ == "__main__": + main() From bfeff88063f0d4f66508423d0e72bc7590360076 Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 12 Mar 2026 12:06:10 +0100 Subject: [PATCH 2/5] no fucking pickle --- .../integrations/pipeline_parallel.py | 52 ++++++++++++++++--- 1 file changed, 46 insertions(+), 6 deletions(-) diff --git a/src/transformers/integrations/pipeline_parallel.py b/src/transformers/integrations/pipeline_parallel.py index 8a5d4455a29b..877c0e50ca7d 100644 --- a/src/transformers/integrations/pipeline_parallel.py +++ b/src/transformers/integrations/pipeline_parallel.py @@ -14,7 +14,6 @@ from __future__ import annotations import math -import pickle from contextlib import contextmanager from typing import Iterable, Sequence @@ -32,6 +31,49 @@ logger = logging.get_logger(__name__) +# Stable dtype ↔ integer mapping used in tensor metadata encoding (no pickle). +_DTYPE_TO_ID: dict = {} +_ID_TO_DTYPE: dict = {} + +def _build_dtype_maps(): + import torch + pairs = [ + (torch.float32, 0), (torch.float16, 1), (torch.bfloat16, 2), + (torch.float64, 3), (torch.int64, 4), (torch.int32, 5), + (torch.int16, 6), (torch.int8, 7), (torch.uint8, 8), (torch.bool, 9), + ] + for dtype, idx in pairs: + _DTYPE_TO_ID[dtype] = idx + _ID_TO_DTYPE[idx] = dtype + +def _encode_tensor_meta(tensors) -> "torch.Tensor": + """Encode shapes/dtypes into a flat int64 tensor — no pickle.""" + if not _DTYPE_TO_ID: + _build_dtype_maps() + # Format: [N, ndim_0, d0_0, ..., dtype_id_0, ndim_1, ...] + import torch + values = [len(tensors)] + for t in tensors: + values.append(len(t.shape)) + values.extend(t.shape) + values.append(_DTYPE_TO_ID[t.dtype]) + return torch.tensor(values, dtype=torch.int64) + + +def _decode_tensor_meta(meta_tensor) -> list: + """Decode shapes/dtypes from a flat int64 tensor.""" + if not _ID_TO_DTYPE: + _build_dtype_maps() + data = meta_tensor.tolist() + n, idx, result = data[0], 1, [] + for _ in range(n): + ndim = data[idx]; idx += 1 + shape = tuple(data[idx : idx + ndim]); idx += ndim + dtype = _ID_TO_DTYPE[data[idx]]; idx += 1 + result.append((shape, dtype)) + return result + + class PipelineParallelSkipped(nn.Module): """Placeholder module used on ranks that do not own a stage. @@ -169,9 +211,7 @@ def send_tensors(tensors: Sequence[torch.Tensor], dst: int, group=None, use_cuda if len(tensors) == 0: return - meta = [(tuple(t.shape), t.dtype) for t in tensors] - meta_bytes = pickle.dumps(meta) - meta_tensor = torch.tensor(list(meta_bytes), dtype=torch.uint8) + meta_tensor = _encode_tensor_meta(tensors) meta_len = torch.tensor([meta_tensor.numel()], dtype=torch.int64) dist.send(meta_len, dst=dst, group=group) dist.send(meta_tensor, dst=dst, group=group) @@ -195,9 +235,9 @@ def recv_tensors(count: int, src: int, device, group=None, use_cuda_stream: bool meta_len = torch.empty(1, dtype=torch.int64) dist.recv(meta_len, src=src, group=group) - meta_tensor = torch.empty(int(meta_len.item()), dtype=torch.uint8) + meta_tensor = torch.empty(int(meta_len.item()), dtype=torch.int64) dist.recv(meta_tensor, src=src, group=group) - meta = pickle.loads(bytes(meta_tensor.tolist())) + meta = _decode_tensor_meta(meta_tensor) tensors: list[torch.Tensor] = [] stream = torch.cuda.Stream() if use_cuda_stream and torch.cuda.is_available() and device is not None and torch.device(device).type == "cuda" else None handles = [] From 94602eb24ac383c7ce075fa3be6c784ee96ec17e Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 12 Mar 2026 12:46:08 +0100 Subject: [PATCH 3/5] simplify ai slope --- .../integrations/pipeline_parallel.py | 166 ++++++------------ 1 file changed, 52 insertions(+), 114 deletions(-) diff --git a/src/transformers/integrations/pipeline_parallel.py b/src/transformers/integrations/pipeline_parallel.py index 877c0e50ca7d..817e75340483 100644 --- a/src/transformers/integrations/pipeline_parallel.py +++ b/src/transformers/integrations/pipeline_parallel.py @@ -13,9 +13,9 @@ # limitations under the License. from __future__ import annotations +import functools import math -from contextlib import contextmanager -from typing import Iterable, Sequence +from typing import Sequence from ..utils import logging from ..utils.import_utils import is_torch_available @@ -91,26 +91,6 @@ def forward(self, *args, **kwargs): return args[0] if len(args) > 0 else None -@contextmanager -def _silent_non_primary(local_rank: int): - """Silence stdout/stderr on non-primary ranks to reduce clutter.""" - - if local_rank == 0: - yield - else: - import os - import sys - - old_stdout, old_stderr = sys.stdout, sys.stderr - sys.stdout = open(os.devnull, "w") - sys.stderr = open(os.devnull, "w") - try: - yield - finally: - sys.stdout.close() - sys.stderr.close() - sys.stdout, sys.stderr = old_stdout, old_stderr - def initialize_pipeline_parallelism(pp_plan: dict[str, tuple[list[str], list[str]]] | None, pp_size: int | None = None, device_mesh=None, device_map=None): """Infer pipeline parallel rank/size and device mapping. @@ -327,6 +307,30 @@ def _pipeline_coordinates(device_mesh): return coord[pp_dim], device_mesh.size(pp_dim) if device_mesh is not None else dist.get_world_size(), tp_rank +def _recv_pre_hook(mod, args, kwargs, *, pp_rank, pp_group, n_inputs, use_cuda_stream): + """Receive activations from the previous pipeline stage.""" + if pp_rank == 0: + return args, kwargs + device = args[0].device if args and hasattr(args[0], "device") else next(mod.parameters()).device + local_rank = dist.get_rank(pp_group) if pp_group is not None else dist.get_rank() + src = dist.get_global_rank(pp_group, local_rank - 1) if pp_group is not None else local_rank - 1 + tensors = [_PipelineRecvFn.apply(1, src, device, pp_group, use_cuda_stream) for _ in range(n_inputs)] + return tuple(tensors), {} + + +def _send_post_hook(mod, args, kwargs, output, *, pp_rank, pp_size, pp_group, n_outputs, use_cuda_stream): + """Send activations to the next pipeline stage.""" + if pp_rank >= pp_size - 1: + return output + outputs = output if isinstance(output, (tuple, list)) else (output,) + local_rank = dist.get_rank(pp_group) if pp_group is not None else dist.get_rank() + dst = dist.get_global_rank(pp_group, local_rank + 1) if pp_group is not None else local_rank + 1 + to_send = [outputs[min(i, len(outputs) - 1)] for i in range(n_outputs)] + for t in to_send: + _PipelineSendFn.apply(t, dst, pp_group, use_cuda_stream) + return output + + def add_pipeline_parallel_hooks( model: nn.Module, pp_plan: dict[str, tuple[list[str], list[str]]], @@ -359,101 +363,35 @@ def add_pipeline_parallel_hooks( partitions = partition_units(units, pp_size) local_unit_names = set(partitions[pp_rank]) if pp_rank < len(partitions) else set() - # Replace non-local modules with skipped placeholders to avoid computation - for name, module, _, _ in units: + local_indices = [i for i, (name, _, _, _) in enumerate(units) if name in local_unit_names] + if not local_indices: + return model + first_local, last_local = local_indices[0], local_indices[-1] + + for i, (name, module, input_names, output_names) in enumerate(units): if name not in local_unit_names: parent_path, child_name = name.rsplit(".", 1) if "." in name else ("", name) - parent = model.get_submodule(parent_path) if parent_path != "" else model + parent = model.get_submodule(parent_path) if parent_path else model setattr(parent, child_name, PipelineParallelSkipped()) - - # Identify first/last local modules to add comm hooks - local_units_ordered = [u for u in units if u[0] in local_unit_names] - if len(local_units_ordered) == 0: - return model - first_name, first_mod, first_inputs, _ = local_units_ordered[0] - last_name, last_mod, _, last_outputs = local_units_ordered[-1] - - runtime_key = "_hf_pp_runtime" - - def _init_runtime(module, args, kwargs): - # args/kwargs are from model forward; seed activation dict with kwargs - runtime = {"activations": {}, "device": args[0].device if len(args) > 0 and hasattr(args[0], "device") else None} - runtime["activations"].update(kwargs) - setattr(model, runtime_key, runtime) - return None - - def _cleanup_runtime(module, args, kwargs, output): - if hasattr(model, runtime_key): - delattr(model, runtime_key) - return output - - model.register_forward_pre_hook(_init_runtime, with_kwargs=True) - model.register_forward_hook(_cleanup_runtime, with_kwargs=True) - - def _recv_pre_hook(mod, args, kwargs): - runtime = getattr(model, runtime_key) - local_pp_rank, global_pp_rank, _ = _group_ranks(pp_group) - if pp_rank > 0: - src_local = local_pp_rank - 1 - src_global = dist.get_global_rank(pp_group, src_local) if pp_group is not None else global_pp_rank - 1 - with _silent_non_primary(pp_rank): - tensors = [ - _PipelineRecvFn.apply( - 1, - src_global, - runtime.get("device"), - pp_group, - use_cuda_stream, - ) - for _ in first_inputs - ] - for name, tensor in zip(first_inputs, tensors): - runtime["activations"][name] = tensor - return tuple(tensors), {} - else: - return args, kwargs - # Build kwargs for this module from runtime activations - new_kwargs = {name: runtime["activations"].get(name) for name in first_inputs if name in runtime["activations"]} - return (), new_kwargs - - def _send_post_hook(mod, args, kwargs, output): - runtime = getattr(model, runtime_key) - outputs = output if isinstance(output, (tuple, list)) else (output,) - to_send = [] - for name in last_outputs: - # Map sequentially to module outputs - idx = last_outputs.index(name) - if idx < len(outputs): - tensor = outputs[idx] - else: - tensor = outputs[0] - runtime["activations"][name] = tensor - to_send.append(tensor) - - if pp_rank < pp_size - 1: - local_pp_rank, global_pp_rank, _ = _group_ranks(pp_group) - dst_local = local_pp_rank + 1 - dst_global = dist.get_global_rank(pp_group, dst_local) if pp_group is not None else global_pp_rank + 1 - to_send = [_PipelineSendFn.apply(tensor, dst_global, pp_group, use_cuda_stream) for tensor in to_send] - return output - - first_mod.register_forward_pre_hook(_recv_pre_hook, with_kwargs=True) - last_mod.register_forward_hook(_send_post_hook, with_kwargs=True) - - # For all local modules, store outputs into activation dict for downstream - def _store_outputs(mod, args, kwargs, output, output_names): - runtime = getattr(model, runtime_key) - outputs = output if isinstance(output, (tuple, list)) else (output,) - for idx, name in enumerate(output_names): - if idx < len(outputs): - runtime["activations"][name] = outputs[idx] - # For first module when not pipeline root, ensure gradients propagate - if pp_rank > 0 and outputs[idx].requires_grad and name in first_inputs: - outputs[idx].register_hook(lambda grad, name=name: grad) - return output - - for name, module, _, out_names in local_units_ordered: - module.register_forward_hook(lambda m, a, kw, o, out_names=out_names: _store_outputs(m, a, kw, o, out_names), with_kwargs=True) + continue + if i == first_local: + module.register_forward_pre_hook( + functools.partial( + _recv_pre_hook, + pp_rank=pp_rank, pp_group=pp_group, + n_inputs=len(input_names), use_cuda_stream=use_cuda_stream, + ), + with_kwargs=True, + ) + if i == last_local: + module.register_forward_hook( + functools.partial( + _send_post_hook, + pp_rank=pp_rank, pp_size=pp_size, pp_group=pp_group, + n_outputs=len(output_names), use_cuda_stream=use_cuda_stream, + ), + with_kwargs=True, + ) return model From 7b1f9e27f92b0769661317c83709e885a38558a3 Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 12 Mar 2026 12:47:35 +0100 Subject: [PATCH 4/5] more --- src/transformers/integrations/pipeline_parallel.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/transformers/integrations/pipeline_parallel.py b/src/transformers/integrations/pipeline_parallel.py index 817e75340483..3a28ec2db4e9 100644 --- a/src/transformers/integrations/pipeline_parallel.py +++ b/src/transformers/integrations/pipeline_parallel.py @@ -363,18 +363,14 @@ def add_pipeline_parallel_hooks( partitions = partition_units(units, pp_size) local_unit_names = set(partitions[pp_rank]) if pp_rank < len(partitions) else set() - local_indices = [i for i, (name, _, _, _) in enumerate(units) if name in local_unit_names] - if not local_indices: - return model - first_local, last_local = local_indices[0], local_indices[-1] - - for i, (name, module, input_names, output_names) in enumerate(units): + local_idx = 0 + for name, module, input_names, output_names in units: if name not in local_unit_names: parent_path, child_name = name.rsplit(".", 1) if "." in name else ("", name) parent = model.get_submodule(parent_path) if parent_path else model setattr(parent, child_name, PipelineParallelSkipped()) continue - if i == first_local: + if local_idx == 0: module.register_forward_pre_hook( functools.partial( _recv_pre_hook, @@ -383,7 +379,7 @@ def add_pipeline_parallel_hooks( ), with_kwargs=True, ) - if i == last_local: + if local_idx == len(local_unit_names) - 1: module.register_forward_hook( functools.partial( _send_post_hook, From 8dbcf0c79f4f47f6210719d8d4c498e340600a49 Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 12 Mar 2026 12:53:17 +0100 Subject: [PATCH 5/5] nit --- .../integrations/pipeline_parallel.py | 40 +++++++++---------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/src/transformers/integrations/pipeline_parallel.py b/src/transformers/integrations/pipeline_parallel.py index 3a28ec2db4e9..10a157497ad7 100644 --- a/src/transformers/integrations/pipeline_parallel.py +++ b/src/transformers/integrations/pipeline_parallel.py @@ -252,38 +252,39 @@ def recv_grads(count: int, src: int, device, group=None, use_cuda_stream: bool = class _PipelineRecvFn(Function): + """Receive N tensors from the previous stage; send N grads back on backward.""" + @staticmethod - def forward(ctx, count: int, src: int, device, group, use_cuda_stream: bool): + def forward(ctx, src: int, device, group, use_cuda_stream: bool, n: int): ctx.src = src ctx.group = group ctx.use_cuda_stream = use_cuda_stream - ctx.device = device - tensors = recv_tensors(count, src=src, device=device, group=group, use_cuda_stream=use_cuda_stream) - if count == 1: - return tensors[0] - return tuple(tensors) + ctx.n = n + return tuple(recv_tensors(n, src=src, device=device, group=group, use_cuda_stream=use_cuda_stream)) @staticmethod def backward(ctx, *grad_outputs): - grads = [g for g in grad_outputs] - send_grads(grads, dst=ctx.src, group=ctx.group, use_cuda_stream=ctx.use_cuda_stream) - # None gradients for non-tensor args + send_grads(list(grad_outputs), dst=ctx.src, group=ctx.group, use_cuda_stream=ctx.use_cuda_stream) return (None, None, None, None, None) class _PipelineSendFn(Function): + """Send N tensors to the next stage; receive N grads back on backward.""" + @staticmethod - def forward(ctx, tensor: torch.Tensor, dst: int, group, use_cuda_stream: bool): + def forward(ctx, dst: int, group, use_cuda_stream: bool, *tensors: torch.Tensor): ctx.dst = dst ctx.group = group ctx.use_cuda_stream = use_cuda_stream - send_tensors([tensor], dst=dst, group=group, use_cuda_stream=use_cuda_stream) - return tensor + ctx.device = tensors[0].device if tensors else None + ctx.n = len(tensors) + send_tensors(tensors, dst=dst, group=group, use_cuda_stream=use_cuda_stream) + return tensors @staticmethod - def backward(ctx, grad_output: torch.Tensor): - grads = recv_grads(1, src=ctx.dst, device=grad_output.device, group=ctx.group, use_cuda_stream=ctx.use_cuda_stream) - return grads[0], None, None, None + def backward(ctx, *grad_outputs): + grads = recv_grads(ctx.n, src=ctx.dst, device=ctx.device, group=ctx.group, use_cuda_stream=ctx.use_cuda_stream) + return (None, None, None) + tuple(grads) def _pipeline_coordinates(device_mesh): @@ -314,8 +315,8 @@ def _recv_pre_hook(mod, args, kwargs, *, pp_rank, pp_group, n_inputs, use_cuda_s device = args[0].device if args and hasattr(args[0], "device") else next(mod.parameters()).device local_rank = dist.get_rank(pp_group) if pp_group is not None else dist.get_rank() src = dist.get_global_rank(pp_group, local_rank - 1) if pp_group is not None else local_rank - 1 - tensors = [_PipelineRecvFn.apply(1, src, device, pp_group, use_cuda_stream) for _ in range(n_inputs)] - return tuple(tensors), {} + tensors = _PipelineRecvFn.apply(src, device, pp_group, use_cuda_stream, n_inputs) + return tensors, {} def _send_post_hook(mod, args, kwargs, output, *, pp_rank, pp_size, pp_group, n_outputs, use_cuda_stream): @@ -325,9 +326,8 @@ def _send_post_hook(mod, args, kwargs, output, *, pp_rank, pp_size, pp_group, n_ outputs = output if isinstance(output, (tuple, list)) else (output,) local_rank = dist.get_rank(pp_group) if pp_group is not None else dist.get_rank() dst = dist.get_global_rank(pp_group, local_rank + 1) if pp_group is not None else local_rank + 1 - to_send = [outputs[min(i, len(outputs) - 1)] for i in range(n_outputs)] - for t in to_send: - _PipelineSendFn.apply(t, dst, pp_group, use_cuda_stream) + to_send = tuple(outputs[min(i, len(outputs) - 1)] for i in range(n_outputs)) + _PipelineSendFn.apply(dst, pp_group, use_cuda_stream, *to_send) return output