# Copyright (c) 2026 PaddlePaddle Authors. 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 collections import OrderedDict from dataclasses import dataclass, field from enum import Enum import numpy as np import paddle import paddle.distributed as dist from paddle.autograd import PyLayer from paddle.distributed.fleet.utils.tensor_fusion_helper import ( align, alignment, get_current_device_type, ) # Global registry for fsdp_context _g_fsdp_context = None def register_fsdp_context(context): global _g_fsdp_context _g_fsdp_context = context def get_fsdp_context(): return _g_fsdp_context class BufferState(Enum): # Buffer status for lazy double buffer mechanism # # State transitions: # FREED ──all_gather──> USING ──computation done──> READY ──release──> FREED # ^ │ # │ (reuse) │ # └────────────────────────────┘ FREED = 1 # Released, buffer data is sharded, tmp_buffer not allocated USING = 2 # Unsharded and actively in use READY = 3 # Unsharded, marked for lazy release, can be reused SYNCING = 4 # Communication in progress @dataclass class BufferGroup: params: list = field(default_factory=list) dtype: object = None trainable: bool = None fsdp_unit_id: int = None is_tie: bool = False is_expert_param: bool = False fsdp_group: object = None params_buffer: 'TensorFusionBuffer' = None grads_buffer: 'TensorFusionBuffer' = None params_use_sum: int = 0 params_use_cnt: int = 0 grads_use_sum: int = 0 grads_use_cnt: int = 0 def _dtensor_from_local(local_tensor, mesh, placements): global_dims = list(local_tensor.shape) for idx, placement in enumerate(placements): if placement.is_shard(): global_dims[placement.get_dim()] = ( global_dims[placement.get_dim()] * mesh.shape[idx] ) place = paddle.framework._current_expected_place() place = paddle.framework._get_paddle_place(place) return paddle.Tensor( local_tensor, dims=global_dims, process_mesh=mesh, placements=placements, place=place, ) class TensorFusionBuffer: def __init__( self, group_id, params, fsdp_degree, dtype, is_params=False, main_grad_dtype=None, ): # Calculate total buffer size needed (with padding) self.group_id = group_id self.fsdp_degree = fsdp_degree self.dtype = dtype self.main_grad_dtype = ( main_grad_dtype if main_grad_dtype is not None else dtype ) self.total_buffer_size = 0 self.param_offsets = {} self.tmp_data_buffer = None self.comm_task = None self.trainable = params[0].trainable for param in params: self.param_offsets[param.name] = self.total_buffer_size self.total_buffer_size += self.get_padded_size(param) if is_params: # Create fused params_buffer # TODO(lizhenxing): Build full params_buffer on CPU and only move shards to GPU to minimize mem peaks self.data_buffer = paddle.zeros( shape=[self.total_buffer_size], dtype=dtype, ) # Use BufferState enum instead of is_shard boolean, initial state is FREED (sharded) self.status = BufferState.FREED for param in params: offset = self.param_offsets[param.name] stop_gradient = param.stop_gradient local_shape = param._local_shape param.stop_gradient = True param._local_value().flatten_() paddle.assign( param._local_value(), self.data_buffer._slice( offset, offset + param._numel(), ), ) param._clear_data() param.stop_gradient = stop_gradient param._local_value().get_tensor()._set_dims(local_shape) paddle.device.cuda.empty_cache() mesh = dist.auto_parallel.get_mesh() curr_global_rank = paddle.distributed.get_rank() if curr_global_rank in mesh.process_ids: total_nums = self.data_buffer.shape[0] num_of_pieces = mesh.shape[0] piece_len = (total_nums + num_of_pieces - 1) // num_of_pieces rank_relative = mesh.process_ids.index(curr_global_rank) start = rank_relative * piece_len end = min(start + piece_len, total_nums) self.data_buffer = paddle.slice( self.data_buffer, [0], [start], [end] ).clone() # Init params_buffer attr self.data_buffer.name = "fuse_params_" + str(group_id) self.data_buffer.stop_gradient = params[0].stop_gradient self.data_buffer.optimize_attr = params[0].optimize_attr else: # Create fused grads_buffer with shard self.data_buffer = paddle.zeros( shape=[self.total_buffer_size // self.fsdp_degree], dtype=self.main_grad_dtype, ) # Register get_main_grad method for each param, returns view_slice of grad_buffer for param in params: if param.trainable: param._fusion_buffer = self param._param_offsets = self.param_offsets def get_grad_from_tmp_buf(param): tmp_buffer = param._fusion_buffer.get_tmp_buffer() offset = param._param_offsets[param.name] main_grad = paddle._C_ops.view_slice( tmp_buffer, offset, offset + param._numel(), ) return main_grad param.get_main_grad = get_grad_from_tmp_buf.__get__(param) def get_padded_size(self, param): size = np.prod(param.shape) align_size = ( alignment[get_current_device_type()] // align[param.dtype] * self.fsdp_degree ) return ((size + align_size - 1) // align_size) * align_size def get_tmp_buffer(self): # Reuse tmp_buffer if exists, else create if self.tmp_data_buffer is None: self.tmp_data_buffer = paddle.zeros( shape=[self.total_buffer_size], dtype=self.dtype ) return self.tmp_data_buffer def clear_tmp_buffer(self): if self.tmp_data_buffer is not None: self.tmp_data_buffer._clear_data() self.tmp_data_buffer = None # paddle.device.cuda.empty_cache() class FSDPBufferManager: def __init__( self, model, mesh, fsdp_unit_layers=None, moe_layers_name=None ): self.model = model self._fsdp_group = mesh.get_group("dp") self.main_grad_dtype = paddle.float32 # Get EP group if "ep" dimension exists in mesh if "ep" in mesh.dim_names: self._ep_fsdp_group = mesh.get_group("ep") else: self._ep_fsdp_group = self._fsdp_group topk = None if hasattr(self.model, 'config') and hasattr( self.model.config, 'num_experts_per_tok' ): topk = self.model.config.num_experts_per_tok # Layer types to wrap as FSDP sharding layers # Note: 'Qwen3VLTextDecoderLayer' is temporary; fleet models all use 'TransformerLayer' self.fsdp_unit_layers = fsdp_unit_layers or [ 'TransformerLayer', 'Qwen3VLTextDecoderLayer', 'Qwen3MoeDecoderLayer', ] # Layer types to identify MoE expert layers self.moe_layers_name = moe_layers_name or [ 'StandardMLPExpert', ] # Get tie_param_name if using tie_weights self.tie_param_name = None # Note: need add get_input_embeddings in fleet modeling # if hasattr(self.model, "get_input_embeddings"): # self.tie_param_name = self.model.get_input_embeddings().weight.name # Create buffer_groups grouped_params, group_is_expert = self._build_groups() self.buffer_groups = [] self.param_to_buffer_id = {} # Create params_buffer, grads_buffer with groups for gid, params in grouped_params.items(): is_expert = group_is_expert.get(gid, False) # Use EP group for expert params, DP group for regular params fsdp_group = self._ep_fsdp_group if is_expert else self._fsdp_group params_buffer = TensorFusionBuffer( gid, params, fsdp_group.nranks, params[0].dtype, is_params=True, ) if not params[0].stop_gradient: grads_buffer = TensorFusionBuffer( gid, params, fsdp_group.nranks, params[0].dtype, main_grad_dtype=self.main_grad_dtype, ) else: grads_buffer = None if is_expert: _params_use_sum = topk _grads_use_sum = topk else: _params_use_sum = len(params) _grads_use_sum = len(params) self.buffer_groups.append( BufferGroup( params=params, dtype=params[0].dtype, trainable=params[0].trainable, is_expert_param=is_expert, fsdp_group=fsdp_group, params_buffer=params_buffer, grads_buffer=grads_buffer, params_use_sum=_params_use_sum, params_use_cnt=0, grads_use_sum=_grads_use_sum, grads_use_cnt=0, ) ) for param in params: self.param_to_buffer_id[param.name] = gid def _build_groups(self): parameters = self.model.parameters() grouped_params = OrderedDict() group_is_expert = {} curr_gid = 0 param_to_unit_id = {} for unit_id, module in enumerate(self.model.modules()): if type(module).__name__ in self.fsdp_unit_layers: for param in module.parameters(): param_to_unit_id[param.name] = unit_id if type(module).__name__ in self.moe_layers_name: for param in module.parameters(): param.is_moe_param = True temp_groups = [] for param in parameters: name = param.name is_expert = getattr(param, "is_moe_param", False) if is_expert: continue is_tie = ( self.tie_param_name is not None and name == self.tie_param_name ) param_attrs = { "dtype": param.dtype, "trainable": param.trainable, "fsdp_unit_id": param_to_unit_id.get(name), "is_tie": is_tie, "is_expert_param": is_expert, } found_group = False for param_group in temp_groups: if ( param_group.dtype == param_attrs["dtype"] and param_group.trainable == param_attrs["trainable"] and param_group.fsdp_unit_id == param_attrs["fsdp_unit_id"] and param_group.is_tie == param_attrs["is_tie"] and param_group.is_expert_param == param_attrs["is_expert_param"] ): param_group.params.append(param) found_group = True break # Create new group if no matching if not found_group: temp_groups.append(BufferGroup(params=[param], **param_attrs)) def group_sort_key(group): priority = 0 if group.is_tie else (1 if not group.trainable else 2) return ( priority, group.fsdp_unit_id if group.fsdp_unit_id is not None else float('inf'), ) sorted_groups = sorted(temp_groups, key=group_sort_key) # For each sorted parameter group, buffer them by execution order for param_group in sorted_groups: cur_params = param_group.params if len(cur_params) == 0: continue for p in cur_params: grouped_params.setdefault(curr_gid, []).append(p) group_is_expert[curr_gid] = param_group.is_expert_param curr_gid += 1 return grouped_params, group_is_expert class FSDPCommManager: def __init__( self, buffer_manager, enable_overlap=True, double_buffer_limit=2, ): self.buffer_manager = buffer_manager self.enable_overlap = enable_overlap self.grad_reduce_queue = [] # for double buffer mechanism config self.double_buffer_limit = double_buffer_limit self.buffer_cnt_in_using = 0 self._need_zero_grads = True def _release_one_buffer_if_needed(self): # Release a buffer with the READY status if needed while self.buffer_cnt_in_using >= self.double_buffer_limit: found = False for group in self.buffer_manager.buffer_groups: if group.params_buffer.status == BufferState.READY: group.params_buffer.status = BufferState.FREED group.params_buffer.clear_tmp_buffer() self.buffer_cnt_in_using -= 1 found = True break if not found: break def _next_buffer_id(self, gid, is_backward): # Get next buffer id for prefetch if is_backward: next_gid = gid - 1 # Search backward for trainable buffer_groups while ( next_gid >= 0 and not self.buffer_manager.buffer_groups[ next_gid ].params_buffer.trainable ): next_gid -= 1 return max(next_gid, 0) else: return min(gid + 1, len(self.buffer_manager.buffer_groups) - 1) def all_gather_params(self, params, is_backward=False): if len(params) == 0: return for param in params: if hasattr(param, "is_moe_param"): continue gid = self.buffer_manager.param_to_buffer_id[param.name] group = self.buffer_manager.buffer_groups[gid] group.params_use_cnt += 1 params_buffer = group.params_buffer # Use group-specific fsdp_group fsdp_group = group.fsdp_group or self.buffer_manager._fsdp_group # Double buffer: reuse buffer if status is READY if params_buffer.status == BufferState.READY: # Reuse: READY -> USING, no need to all_gather again params_buffer.status = BufferState.USING # Overlap prefetch comm if self.enable_overlap: next_gid = self._next_buffer_id(gid, is_backward) next_group = self.buffer_manager.buffer_groups[next_gid] next_params_buffer = next_group.params_buffer next_fsdp_group = ( next_group.fsdp_group or self.buffer_manager._fsdp_group ) if next_params_buffer.status == BufferState.FREED: # Check double_buffer_limit before prefetch self._release_one_buffer_if_needed() next_params_buffer.status = BufferState.SYNCING tmp_buffer_prefetch = next_params_buffer.get_tmp_buffer() next_params_buffer.comm_task = ( paddle.distributed.all_gather( tmp_buffer_prefetch, next_params_buffer.data_buffer, group=next_fsdp_group, sync_op=False, ) ) self.buffer_cnt_in_using += 1 # Wait for async comm to complete: SYNCING -> USING if params_buffer.status == BufferState.SYNCING: params_buffer.status = BufferState.USING params_buffer.comm_task.wait() params_buffer.comm_task = None tmp_buffer = params_buffer.get_tmp_buffer() # Do all_gather in sync: FREED -> USING if params_buffer.status == BufferState.FREED: fsdp_group.process_group.all_gather( params_buffer.data_buffer, tmp_buffer ).wait() params_buffer.status = BufferState.USING self.buffer_cnt_in_using += 1 # Bind the unsharded param to the real param offset = params_buffer.param_offsets[param.name] tmp_param = paddle._C_ops.view_slice( tmp_buffer, offset, offset + param._numel(), ) tmp_param.get_tensor()._set_dims(param.shape) tmp_param = _dtensor_from_local( tmp_param, param.process_mesh, param.placements, ) param.get_tensor()._share_data_with(tmp_param.get_tensor()) def shard_params(self, params, is_backward=False): affected_gids = set() for param in params: if hasattr(param, "is_moe_param"): continue gid = self.buffer_manager.param_to_buffer_id.get(param.name) group = self.buffer_manager.buffer_groups[gid] stop_gradient = param.stop_gradient local_shape = param._local_shape param._clear_data() param.stop_gradient = stop_gradient param._local_value().get_tensor()._set_dims(local_shape) affected_gids.add(gid) for gid in affected_gids: group = self.buffer_manager.buffer_groups[gid] if group.params_buffer.status == BufferState.USING: group.params_buffer.status = BufferState.READY def reduce_scatter_grads(self, param): if self._need_zero_grads: self._need_zero_grads = False for group in self.buffer_manager.buffer_groups: if group.grads_buffer is not None: group.grads_buffer.data_buffer.zero_() gid = self.buffer_manager.param_to_buffer_id.get(param.name) group = self.buffer_manager.buffer_groups[gid] group.grads_use_cnt += 1 fsdp_group = group.fsdp_group or self.buffer_manager._fsdp_group param.main_grad = None if group.grads_use_cnt == group.grads_use_sum: group.grads_use_cnt = 0 # reduce_scatter from tmp_grad_buffer into grads_buffer grads_buffer = group.grads_buffer # Grad queue mechanism: wait and release completed reduce_scatter async tasks self._wait_for_grad_comm() tmp_buffer = grads_buffer.get_tmp_buffer() shard_size = grads_buffer.data_buffer.shape[0] grad_buffer_shard = tmp_buffer._slice(0, shard_size) if self.enable_overlap: # Comm grads async and check all comm_task before optimizer update grads_buffer.comm_task = paddle.distributed.reduce_scatter( grad_buffer_shard, tmp_buffer, op=paddle.distributed.ReduceOp.SUM, group=fsdp_group, sync_op=False, ) # Add async task to queue self.grad_reduce_queue.append(grads_buffer) else: paddle.distributed.reduce_scatter( grad_buffer_shard, tmp_buffer, op=paddle.distributed.ReduceOp.SUM, group=fsdp_group, sync_op=False, ).wait() grads_buffer.data_buffer.add_(grad_buffer_shard) grads_buffer.clear_tmp_buffer() def _wait_for_grad_comm(self, queue_limit=2): # Wait for async reduce_scatter tasks to complete and release resources # queue_limit: max queue size, default use 2, 0 means wait for all while len(self.grad_reduce_queue) > queue_limit: grads_buffer = self.grad_reduce_queue.pop(0) if grads_buffer.comm_task is not None: grads_buffer.comm_task.wait() grads_buffer.comm_task = None tmp_buffer = grads_buffer.get_tmp_buffer() shard_size = grads_buffer.data_buffer.shape[0] grad_buffer_shard = tmp_buffer._slice(0, shard_size) grads_buffer.data_buffer.add_(grad_buffer_shard) grads_buffer.clear_tmp_buffer() def _finish_grads_sync(self): # Wait for all async reduce_scatter tasks, call before optimizer.step() self._wait_for_grad_comm(queue_limit=0) def _reset_params_buffer_status(self): for group in self.buffer_manager.buffer_groups: params_buffer = group.params_buffer if params_buffer.status in (BufferState.READY, BufferState.USING): # Clear stale tmp_buffer to force re-all_gather with updated data_buffer params_buffer.clear_tmp_buffer() params_buffer.status = BufferState.FREED if self.buffer_cnt_in_using > 0: self.buffer_cnt_in_using -= 1 class FusionBackwardHook(PyLayer): @staticmethod def forward(ctx, *inputs, layer, comm_manager, recursive=False): ctx.layer = layer ctx.comm_manager = comm_manager ctx.recursive = recursive return inputs if len(inputs) > 1 else inputs[0] @staticmethod def backward(ctx, *args): trainable_params = [] for param in ctx.layer.parameters(include_sublayers=ctx.recursive): if param.trainable: trainable_params.append(param) ctx.comm_manager.all_gather_params(trainable_params, is_backward=True) return args class FusionForwardHook(PyLayer): @staticmethod def forward(ctx, *inputs, layer, comm_manager, recursive=False): ctx.layer = layer ctx.comm_manager = comm_manager ctx.recursive = recursive return inputs @staticmethod def backward(ctx, *args): ctx.comm_manager.shard_params( ctx.layer.parameters(include_sublayers=ctx.recursive), is_backward=True, ) return args class FullyShardFusion: def __init__( self, model, mesh, fsdp_unit_layers=None, moe_layers_name=None ): self.model = model self.mesh = self._check_mesh(mesh) self._shard_all_params() self.buffer_manager = FSDPBufferManager( self.model, self.mesh, fsdp_unit_layers, moe_layers_name ) self.comm_manager = FSDPCommManager(self.buffer_manager) self.register_tensor_fusion_hooks(self.model) register_fsdp_context(self) def _check_mesh(self, mesh, pp_idx=0): if "pp" in mesh.dim_names: mesh = mesh.get_mesh_with_dim("pp", pp_idx) return mesh def _shard_all_params(self): def shard_layer_param(layer): for param_name in list(layer._parameters.keys()): param = getattr(layer, param_name) if param is not None: param_placements = [ dist.Replicate() for _ in range(len(self.mesh.shape)) ] if not param.is_dist(): param = dist.shard_tensor( param, self.mesh, param_placements ) setattr(layer, param_name, param) for name, layer in self.model.named_sublayers(include_self=True): shard_layer_param(layer) def comm_sync_and_reset_status(self): self.comm_manager._finish_grads_sync() self.comm_manager._reset_params_buffer_status() self.comm_manager._need_zero_grads = True # Reset main_grad for all trainable parameters for param in self.model.parameters(): if param.trainable: param.main_grad = None def register_tensor_fusion_hooks(self, model): def _pre_forward_hook(sublayers, recursive=False): comm_manager = self.comm_manager @paddle.autograd.no_grad() def all_gather_comm(*_): comm_manager.all_gather_params( sublayers.parameters(include_sublayers=recursive) ) return all_gather_comm def _post_forward_hook(sublayers, recursive=False): comm_manager = self.comm_manager @paddle.autograd.no_grad() def shard_comm(*_): comm_manager.shard_params( sublayers.parameters(include_sublayers=recursive) ) return shard_comm def _update_main_grad_hook(param): comm_manager = self.comm_manager @paddle.autograd.no_grad() def comm_hook(grad): if grad is not None and grad._is_initialized(): # Share mem with grads_tmp_buffer _main_grad = param.get_main_grad() _main_grad.get_tensor()._set_dims(grad._local_shape) param.main_grad = _dtensor_from_local( _main_grad, grad.process_mesh, grad.placements, ) param.main_grad._local_value().copy_(grad._local_value()) grad._clear_data() comm_manager.shard_params([param], is_backward=True) comm_manager.reduce_scatter_grads(param) return comm_hook def _post_backward_hook(param): param.main_grad = None if hasattr(param, "get_main_grad"): param._register_grad_hook(_update_main_grad_hook(param)) for param in model.parameters(): if param.trainable: _post_backward_hook(param) def _register_recursive(layer): is_unit = ( type(layer).__name__ in self.buffer_manager.fsdp_unit_layers ) if is_unit: # For FSDP Unit, register recursive hooks and stop recursion layer.register_forward_pre_hook( _pre_forward_hook(layer, recursive=True) ) layer.register_forward_post_hook( _post_forward_hook(layer, recursive=True) ) self._register_fusion_layer_hooks(layer, recursive=True) return if layer.parameters(include_sublayers=False): layer.register_forward_pre_hook( _pre_forward_hook(layer, recursive=False) ) layer.register_forward_post_hook( _post_forward_hook(layer, recursive=False) ) self._register_fusion_layer_hooks(layer, recursive=False) for child in layer.children(): _register_recursive(child) _register_recursive(model) def _register_fusion_layer_hooks(self, layer, recursive=False): def _forward_post_hook(layer, inputs, outputs): if isinstance(outputs, dict): for key, value in outputs.items(): if ( isinstance(value, paddle.Tensor) and not value.stop_gradient ): outputs[key] = FusionBackwardHook.apply( value, layer=layer, comm_manager=self.comm_manager, recursive=recursive, ) return outputs elif isinstance(outputs, tuple): result = FusionBackwardHook.apply( *outputs, layer=layer, comm_manager=self.comm_manager, recursive=recursive, ) if not isinstance(result, tuple): result = (result,) return result else: return FusionBackwardHook.apply( outputs, layer=layer, comm_manager=self.comm_manager, recursive=recursive, ) def _forward_pre_hook(layer, inputs): return FusionForwardHook.apply( *inputs, layer=layer, comm_manager=self.comm_manager, recursive=recursive, ) layer.register_forward_post_hook(_forward_post_hook) # Register an additional hook for tie_weights shard_params for param in layer.parameters(include_sublayers=False): if param.name == self.comm_manager.buffer_manager.tie_param_name: layer.register_forward_pre_hook(_forward_pre_hook)