wehub-resource-sync
1511 lines
63 KiB
Python
1511 lines
63 KiB
Python
# Copyright (c) 2025 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.
|
|
|
|
"""
|
|
MuonShardingOptimizer (Sharding Stage1 V3): Hybrid Tensor-wise + Element-wise Sharding
|
|
==================================================================
|
|
|
|
Designed for Muon optimizer compatibility:
|
|
- 2D (Muon) parameters: assigned as *whole tensors* to ranks (like V1).
|
|
This avoids the expensive sharding gather in Muon's _muon_update.
|
|
- Non-2D (AdamW) parameters: split element-wise via reduce-scatter (like V2).
|
|
This provides memory balancing across ranks.
|
|
|
|
The key insight is that Muon requires the full 2D matrix for Newton-Schulz
|
|
orthogonalisation, so keeping 2D params whole on each rank eliminates the
|
|
need for gather_varlen communication during the optimizer step.
|
|
|
|
Parameters are grouped by their `color` attribute, which specifies the
|
|
communication group to use:
|
|
- color=None or -1: default sharding_group
|
|
- color={'color': <key>, 'group': <group>}: custom group read directly
|
|
from the param, no code changes needed to add new color groups.
|
|
"""
|
|
|
|
import math
|
|
import os
|
|
import warnings
|
|
from collections import defaultdict
|
|
from functools import reduce as functools_reduce
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle import framework
|
|
from paddle.base.framework import EagerParamBase
|
|
from paddle.distributed import fleet
|
|
from paddle.distributed.communication.batch_isend_irecv import (
|
|
_coalescing_manager,
|
|
)
|
|
from paddle.distributed.communication.reduce import (
|
|
ReduceOp,
|
|
is_avg_reduce_op_supported,
|
|
)
|
|
from paddle.distributed.fleet.utils import timer_helper as timer
|
|
from paddle.distributed.fleet.utils.log_util import logger
|
|
from paddle.distributed.fleet.utils.tensor_fusion_helper import (
|
|
HOOK_ACTION,
|
|
FusedCommBuffer,
|
|
assign_group_by_size,
|
|
)
|
|
|
|
g_shard_bypass_dygraph_optimizer = int(
|
|
os.environ.get("FLAGS_shard_bypass_dygraph_optimizer", 0)
|
|
)
|
|
|
|
|
|
def _is_trainable(param):
|
|
return not param.stop_gradient
|
|
|
|
|
|
def get_same_card_rank(moe_sharding_group_ranks, rank):
|
|
"""
|
|
Get the MoE sharding group rank within the same trainer as the specified rank.
|
|
|
|
Args:
|
|
moe_sharding_group_ranks (set): Set of ranks in the MoE sharding group
|
|
rank (int): Current rank
|
|
|
|
Returns:
|
|
int: The rank within the same trainer that belongs to the MoE sharding group;
|
|
returns -1 if not found
|
|
"""
|
|
trainer = rank // 8
|
|
for i in range(trainer * 8, (trainer + 1) * 8):
|
|
if i in moe_sharding_group_ranks:
|
|
return i
|
|
return -1
|
|
|
|
|
|
def get_trainer_ranks(rank):
|
|
"""
|
|
Get the trainer ID and all ranks belonging to that trainer.
|
|
|
|
In distributed training, every 8 GPUs form a "trainer". This function computes:
|
|
1. The trainer number for the given rank.
|
|
2. The range of all 8 ranks within that same trainer.
|
|
|
|
Args:
|
|
rank (int): The global rank ID of the current process
|
|
|
|
Returns:
|
|
tuple:
|
|
- train_id (int): The trainer index (0, 1, 2...)
|
|
- Ranks iterable: All 8 ranks in this trainer,
|
|
e.g., if rank=10 returns (1, range(8, 16))
|
|
|
|
Example:
|
|
>>> get_trainer_ranks(5)
|
|
(0, range(0, 8))
|
|
>>> get_trainer_ranks(12)
|
|
(1, range(8, 16))
|
|
"""
|
|
trainer = rank // 8
|
|
return trainer, range(trainer * 8, (trainer + 1) * 8)
|
|
|
|
|
|
class MuonShardingOptimizer:
|
|
"""
|
|
Hybrid sharding optimizer for Muon:
|
|
- 2D (Muon) parameters: tensor-wise assignment to ranks (no cross-rank split).
|
|
Gradient communication uses reduce; parameter sync uses broadcast.
|
|
- Non-2D (AdamW) parameters: element-wise split across ranks (like V2).
|
|
Gradient communication uses reduce-scatter; parameter sync uses all-gather.
|
|
|
|
Parameters are grouped by `color` attribute to determine the communication
|
|
group. Each color group has its own 2D parameter partition and communication.
|
|
|
|
This avoids the expensive gather_varlen in Muon's _muon_update while
|
|
maintaining memory balance across ranks.
|
|
"""
|
|
|
|
def __init__(self, optimizer, hcg=None):
|
|
logger.info("init MuonShardingOptimizer")
|
|
if isinstance(optimizer._parameter_list[0], dict):
|
|
raise TypeError(
|
|
"Do not support param_groups now, please set optimizer._parameter_list as a list of Parameter"
|
|
)
|
|
if not hasattr(optimizer, '_apply_optimize') or not callable(
|
|
optimizer._apply_optimize
|
|
):
|
|
raise ValueError(
|
|
"the optimizer object should have _apply_optimize function"
|
|
)
|
|
|
|
self._inner_opt = optimizer
|
|
# Get hcg from fleet if not provided
|
|
if hcg is None:
|
|
hcg = fleet.fleet._hcg
|
|
self._hcg = hcg
|
|
self._sharding_world_size = self._hcg.get_sharding_parallel_world_size()
|
|
self._sharding_rank = self._hcg.get_sharding_parallel_rank()
|
|
self._global_rank = paddle.distributed.get_rank()
|
|
|
|
# Temporarily: TP is not supported in MuonShardingOptimizer
|
|
_tp_world_size = self._hcg.get_model_parallel_world_size()
|
|
assert _tp_world_size == 1, (
|
|
f"MuonShardingOptimizer does not support tensor parallelism yet. "
|
|
f"Got tp_world_size={_tp_world_size}. Please set tensor_parallel_degree=1."
|
|
)
|
|
|
|
strategy = fleet.fleet._user_defined_strategy
|
|
sharding_configs = strategy.hybrid_configs['sharding_configs']
|
|
|
|
self.tensor_fusion = sharding_configs.tensor_fusion
|
|
self.accumulate_steps = sharding_configs.accumulate_steps
|
|
self.comm_overlap = sharding_configs.comm_overlap
|
|
self.comm_buffer_size_MB = sharding_configs.comm_buffer_size_MB
|
|
self.use_reduce_avg = sharding_configs.use_reduce_avg
|
|
self.enable_fuse_optimizer_states = (
|
|
sharding_configs.enable_fuse_optimizer_states
|
|
)
|
|
|
|
self.use_group_call_opt = sharding_configs.comm_group_call_opt
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
ep_degree = (
|
|
hcg.get_expert_parallel_world_size()
|
|
if hasattr(hcg, "get_expert_parallel_world_size")
|
|
else 1
|
|
)
|
|
moe_sharding_degree = hcg.get_moe_sharding_parallel_world_size()
|
|
if self.use_group_call_opt:
|
|
assert ep_degree == 8 and moe_sharding_degree != 1, (
|
|
"comm_group_call_opt should be enabled when ep_degree is 8 and moe_sharding_degree is not 1"
|
|
)
|
|
|
|
if self.enable_fuse_optimizer_states:
|
|
self._inner_opt.use_fusion_storage()
|
|
|
|
if self.use_reduce_avg and (not is_avg_reduce_op_supported()):
|
|
self.use_reduce_avg = False
|
|
warnings.warn(
|
|
"nccl reduce_avg requires paddle compiled with cuda and nccl>=2.10.0, "
|
|
"please check compilation setups."
|
|
)
|
|
|
|
pp_overlap = strategy.hybrid_configs['pp_configs'].sharding_comm_overlap
|
|
self.pp_overlap = pp_overlap
|
|
assert not self.pp_overlap, (
|
|
"muon_sharding_optimizer do not support PP overlap"
|
|
)
|
|
|
|
self._use_main_grad = hasattr(optimizer._parameter_list[0], "main_grad")
|
|
|
|
# The full original parameter list
|
|
self._parameter_list = list(optimizer._parameter_list)
|
|
self._origin_parameter_list = list(optimizer._parameter_list)
|
|
|
|
# Build color -> group_info mapping dynamically from param.color attributes
|
|
sharding_group = hcg.get_sharding_parallel_group()
|
|
self._color_to_group_info = self._build_color_to_group_info_from_params(
|
|
self._parameter_list, sharding_group
|
|
)
|
|
|
|
# Get muon_param_info_map from the inner Muon optimizer.
|
|
# Each entry has use_muon=True/False, set by the Trainer before construction.
|
|
self._muon_param_info_map = getattr(
|
|
optimizer, '_muon_param_info_map', {}
|
|
)
|
|
|
|
# ---- Step 1: Separate params into 2D (Muon) and 1D (AdamW) by color ----
|
|
# Parameters are grouped by their `color` attribute:
|
|
# - color=None or -1: default sharding_group (key: None)
|
|
# - color={'color': <key>, 'group': <group>}: custom comm group
|
|
#
|
|
# For each color group:
|
|
# - 2D (Muon) params: whole tensor, assigned to ranks via tensor-wise partition
|
|
# - non-2D (AdamW) params: element-wise split via FusedCommBuffer
|
|
self._params_2d_by_color = defaultdict(
|
|
list
|
|
) # color_key -> list of 2D params
|
|
self._params_1d = [] # all non-2D params
|
|
self.clear_color = set()
|
|
self._color_to_comm_buffer_list = {}
|
|
for p in self._parameter_list:
|
|
if not _is_trainable(p):
|
|
continue
|
|
|
|
# Extract color value
|
|
color = getattr(p, 'color', -1)
|
|
if isinstance(color, dict):
|
|
color_val = color.get('color', -1)
|
|
else:
|
|
color_val = color
|
|
|
|
# Normalize color: treat None/-1 as default (None key)
|
|
if color_val == -1 or color_val is None:
|
|
color_key = None
|
|
else:
|
|
color_key = color_val
|
|
|
|
param_info = self._muon_param_info_map.get(p.name)
|
|
assert param_info is not None, (
|
|
f"Parameter {p.name!r} (shape={list(p.shape)}) has no muon_param_info. "
|
|
f"Trainer._build_muon_param_info_map must set muon_param_info on all "
|
|
f"trainable parameters before MuonShardingOptimizer is constructed."
|
|
)
|
|
use_muon = param_info.use_muon
|
|
|
|
if use_muon:
|
|
self._params_2d_by_color[color_key].append(p)
|
|
else:
|
|
# Non-2D params use element-wise split via FusedCommBuffer
|
|
self._params_1d.append(p)
|
|
|
|
# ---- Step 2: Partition 2D params for each color group ----
|
|
# For each color, compute rank-to-params and param-to-rank mappings
|
|
self._rank2params_2d_by_color = {} # color -> {rank -> [params]}
|
|
self._param2rank_2d_by_color = {} # color -> {param_name -> rank}
|
|
|
|
for color_key, params_2d in self._params_2d_by_color.items():
|
|
group_info = self._color_to_group_info.get(color_key, {})
|
|
world_size = group_info.get('world_size', 1)
|
|
|
|
if world_size <= 1:
|
|
# No partition needed, all params stay on rank 0
|
|
self._rank2params_2d_by_color[color_key] = {0: list(params_2d)}
|
|
self._param2rank_2d_by_color[color_key] = {
|
|
p.name: 0 for p in params_2d
|
|
}
|
|
else:
|
|
# Greedy partition across ranks
|
|
label = color_key if color_key else "default"
|
|
self._rank2params_2d_by_color[color_key] = (
|
|
self._partition_2d_parameters(
|
|
list(params_2d), world_size, label=label
|
|
)
|
|
)
|
|
self._param2rank_2d_by_color[color_key] = {}
|
|
for rank, params in self._rank2params_2d_by_color[
|
|
color_key
|
|
].items():
|
|
for p in params:
|
|
self._param2rank_2d_by_color[color_key][p.name] = rank
|
|
|
|
# Sort params within each color by owner rank for deterministic ordering
|
|
for color_key, params_2d in self._params_2d_by_color.items():
|
|
params_2d.sort(
|
|
key=lambda p: self._param2rank_2d_by_color[color_key][p.name]
|
|
)
|
|
|
|
# 2D params owned by this sharding rank
|
|
self._local_2d = []
|
|
for color_key, params_2d in self._params_2d_by_color.items():
|
|
rank2params_2d_by_color = self._rank2params_2d_by_color[color_key]
|
|
|
|
group_info = self._color_to_group_info[color_key]
|
|
sharding_rank = max(group_info['rank'], 0)
|
|
|
|
self._local_2d.extend(rank2params_2d_by_color[sharding_rank])
|
|
|
|
self.sd_release_grads = (
|
|
strategy.hybrid_configs['pp_configs'].release_gradients
|
|
or sharding_configs.release_gradients
|
|
)
|
|
self._use_fuse_gradients = self.comm_buffer_size_MB > 0
|
|
# ---- Build comm buffers for 2D params (V1-style) ----
|
|
if self._use_fuse_gradients:
|
|
self.comm_buffer_2d = self._build_2d_comm_buffers()
|
|
self.comm_buffer_2d.sort(key=lambda x: (x._comm_group._id, x._dst))
|
|
|
|
# ---- Step 3: Build comm buffers for 1D params (V2-style) ----
|
|
self._slice_params = {}
|
|
self._comm_buffer_list = []
|
|
self._local_parameter_list_1d = [
|
|
self._create_slice_param(p) for p in self._params_1d
|
|
]
|
|
|
|
self.param2bucket = {}
|
|
|
|
self._build_1d_comm_buffers()
|
|
|
|
# ---- Step 4: Build the optimizer's parameter list ----
|
|
# The optimizer should see:
|
|
# - All 2D params assigned to this rank (all colors, as whole tensors)
|
|
# - 1D slice_params for all non-2D params (element-wise shards)
|
|
local_opt_params = list(self._local_2d) + list(
|
|
self._local_parameter_list_1d
|
|
)
|
|
|
|
self._set_inner_opt_attr('_parameter_list', local_opt_params)
|
|
self._set_inner_opt_attr('_param_groups', local_opt_params)
|
|
|
|
# For external iteration (clear_grad, etc.), expose all params
|
|
self._local_parameter_list = local_opt_params
|
|
|
|
self._enable_timer = strategy.hybrid_configs.get(
|
|
"enable_optimizer_timer", False
|
|
)
|
|
if self._enable_timer:
|
|
if not timer.is_timer_initialized():
|
|
timer.set_timers()
|
|
self.timers = timer.get_timers()
|
|
|
|
# --- Per-rank parameter size summary (for load balancing diagnostics) ---
|
|
_sg_group = hcg.get_sharding_parallel_group()
|
|
_N = self._sharding_world_size
|
|
|
|
# 2D params owned by this sharding rank
|
|
_local_2d_numel = sum(
|
|
int(functools_reduce(lambda x, y: x * y, p.shape, 1))
|
|
for p in self._local_2d
|
|
)
|
|
# 1D (AdamW) slice: each rank holds ceil(numel / sharding_world_size) elements.
|
|
_local_1d_numel = sum(
|
|
math.ceil(
|
|
int(functools_reduce(lambda x, y: x * y, p.shape, 1)) / _N
|
|
)
|
|
for p in self._params_1d
|
|
)
|
|
|
|
_local_total_numel = _local_2d_numel + _local_1d_numel
|
|
_local_total_MB = (
|
|
_local_total_numel * 2 / (1024 * 1024)
|
|
) # bf16/fp16 = 2 bytes
|
|
|
|
# All-gather total numel from all sharding ranks in this PP stage
|
|
_local_numel_tensor = paddle.to_tensor(
|
|
[_local_total_numel], dtype='int64'
|
|
)
|
|
_all_numel_list = []
|
|
paddle.distributed.all_gather(
|
|
_all_numel_list, _local_numel_tensor, group=_sg_group
|
|
)
|
|
_all_numel = [int(t.item()) for t in _all_numel_list]
|
|
_all_MB = [n * 2 / (1024 * 1024) for n in _all_numel]
|
|
|
|
_max_MB = max(_all_MB)
|
|
_min_MB = min(_all_MB)
|
|
_imbalance = (_max_MB - _min_MB) / _max_MB if _max_MB > 0 else 0.0
|
|
|
|
if self._sharding_rank == 0:
|
|
logger.info(
|
|
f"[MuonSharding global_rank={self._global_rank} sharding_rank={self._sharding_rank}] "
|
|
f"SliceSize sharding_group ranks={_sg_group.ranks} | "
|
|
f"per-rank MB: {[f'{mb:.1f}' for mb in _all_MB]} | "
|
|
f"max memory diff={_imbalance * 100:.2f}%"
|
|
)
|
|
if self.use_group_call_opt:
|
|
self.trainer_comms = {}
|
|
world_size = paddle.distributed.get_world_size()
|
|
num_trainers = world_size // 8
|
|
for i in range(num_trainers):
|
|
ranks = range(i * 8, (i + 1) * 8)
|
|
group = paddle.distributed.new_group(ranks)
|
|
self.trainer_comms[i] = group
|
|
|
|
# ------------------------------------------------------------------
|
|
# 2D partition (V1-style greedy)
|
|
# ------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _build_color_to_group_info_from_params(parameter_list, default_group):
|
|
"""Build color->group_info mapping dynamically from param.color attributes.
|
|
|
|
When param.color is a dict containing 'group', the comm group is read
|
|
directly from the param — no hcg-specific method registration required.
|
|
|
|
Returns:
|
|
dict: {
|
|
None: {'group': default_group, 'world_size': N, 'rank': r},
|
|
'<key>': {'group': group, 'world_size': M, 'rank': s},
|
|
# additional entries auto-populated from param.color dicts
|
|
}
|
|
"""
|
|
color_to_info = {
|
|
None: {
|
|
'group': default_group,
|
|
'world_size': len(default_group.ranks) if default_group else 1,
|
|
'rank': default_group.rank if default_group else 0,
|
|
}
|
|
}
|
|
for p in parameter_list:
|
|
color = getattr(p, 'color', -1)
|
|
if isinstance(color, dict):
|
|
color_key = color.get('color', -1)
|
|
if (
|
|
color_key not in (-1, None)
|
|
and color_key not in color_to_info
|
|
):
|
|
group = color.get('group', default_group)
|
|
color_to_info[color_key] = {
|
|
'group': group,
|
|
'world_size': len(group.ranks) if group else 1,
|
|
'rank': group.rank if group else 0,
|
|
}
|
|
return color_to_info
|
|
|
|
def _partition_2d_parameters(self, params, world_size, label=""):
|
|
"""Partition 2D parameters among ranks using greedy bin-packing.
|
|
|
|
Only assign parameters to the first n ranks such that total size > comm_buffer_size_MB.
|
|
Remaining ranks get no parameters.
|
|
"""
|
|
mapping = {}
|
|
for rank in range(world_size):
|
|
mapping[rank] = []
|
|
|
|
parameters = list(params)
|
|
parameters.sort(
|
|
key=lambda p: functools_reduce(lambda x, y: x * y, p.shape),
|
|
reverse=True,
|
|
)
|
|
|
|
total_numel = sum(
|
|
functools_reduce(lambda x, y: x * y, p.shape, 1) for p in parameters
|
|
)
|
|
total_size_bytes = total_numel * 4
|
|
total_size_mb = total_size_bytes / (1024**2)
|
|
|
|
buffer_size_mb = (
|
|
self.comm_buffer_size_MB if self.comm_buffer_size_MB > 0 else 256
|
|
)
|
|
min_active_ranks = 1
|
|
if total_size_mb > 0:
|
|
min_active_ranks = max(1, int(total_size_mb / buffer_size_mb) + 1)
|
|
|
|
active_ranks = min(min_active_ranks, world_size)
|
|
sizes = [0] * active_ranks
|
|
|
|
for param in parameters:
|
|
rank = sizes.index(min(sizes))
|
|
mapping[rank].append(param)
|
|
numel = functools_reduce(lambda x, y: x * y, param.shape, 1)
|
|
sizes[rank] += numel
|
|
|
|
return mapping
|
|
|
|
def _build_2d_comm_buffers(self):
|
|
"""Build communication buffers for 2D (Tensor-wise) parameters using all-reduce."""
|
|
group_size = (
|
|
self.comm_buffer_size_MB * 1024 * 1024
|
|
if self.comm_buffer_size_MB > 0
|
|
else 256 * 1024 * 1024
|
|
)
|
|
comm_buffers = []
|
|
|
|
for color_key, params_2d in self._params_2d_by_color.items():
|
|
group_info = self._color_to_group_info.get(color_key, {})
|
|
comm_group = group_info.get('group', None)
|
|
|
|
fused_parameter_group = defaultdict(list)
|
|
|
|
for p in params_2d:
|
|
dst_rank = self._param2rank_2d_by_color[color_key][p.name]
|
|
fused_parameter_group[dst_rank].append(p)
|
|
|
|
absolute_dst_ranks = {
|
|
rank: comm_group.ranks[rank] for rank in fused_parameter_group
|
|
}
|
|
|
|
for dst, params in fused_parameter_group.items():
|
|
var_groups = assign_group_by_size(params, group_size)
|
|
abs_dst = absolute_dst_ranks[dst]
|
|
|
|
buffer = [
|
|
FusedCommBuffer(
|
|
group_idx,
|
|
parameters,
|
|
comm_group,
|
|
self.accumulate_steps,
|
|
act=HOOK_ACTION.REDUCE,
|
|
dst=abs_dst,
|
|
release_grads=self.sd_release_grads,
|
|
use_reduce_avg=True,
|
|
)
|
|
for group_idx, parameters in var_groups.items()
|
|
]
|
|
comm_buffers.extend(buffer)
|
|
|
|
return comm_buffers
|
|
|
|
# ------------------------------------------------------------------
|
|
# 1D slice creation (V2-style)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _create_slice_param(self, param):
|
|
"""Create a placeholder slice parameter for 1D (element-wise) sharding."""
|
|
slice_param = EagerParamBase(shape=[1], dtype=param.dtype)
|
|
slice_param.name = param.name
|
|
|
|
def copy_attr(attr_name):
|
|
if hasattr(param, attr_name):
|
|
setattr(slice_param, attr_name, getattr(param, attr_name))
|
|
|
|
copy_attr("is_distributed")
|
|
copy_attr("optimize_attr")
|
|
copy_attr("do_model_average")
|
|
copy_attr("need_clip")
|
|
copy_attr("no_sync")
|
|
copy_attr("is_firstly_shared")
|
|
|
|
self._slice_params[param.name] = slice_param
|
|
return slice_param
|
|
|
|
def _build_1d_comm_buffers(self):
|
|
"""Build communication buffers for 1D (AdamW) parameters using reduce-scatter."""
|
|
if self.pp_overlap:
|
|
return
|
|
|
|
comm_group = self._hcg.get_sharding_parallel_group()
|
|
group_size = (
|
|
self.comm_buffer_size_MB * 1024 * 1024
|
|
if self.comm_buffer_size_MB > 0
|
|
else 256 * 1024 * 1024
|
|
)
|
|
|
|
# Group 1D params by (color, comm_group) so each group uses its own FusedCommBuffer
|
|
color_dict = defaultdict(list)
|
|
for param in self._params_1d:
|
|
color = getattr(param, 'color', -1)
|
|
color_group = comm_group
|
|
if isinstance(color, dict):
|
|
color_color = color.get('color', -1)
|
|
color_group = color.get('group', comm_group)
|
|
else:
|
|
color_color = color
|
|
color_dict[(color_color, color_group)].append(param)
|
|
|
|
if not self.comm_overlap:
|
|
for color, params in color_dict.items():
|
|
params.sort(key=lambda x: str(x.dtype))
|
|
|
|
group_idx = 0
|
|
for color, params in color_dict.items():
|
|
g_color = color[0]
|
|
g_group = color[1]
|
|
var_groups = assign_group_by_size(params, group_size)
|
|
for _, parameters in var_groups.items():
|
|
buffer = FusedCommBuffer(
|
|
group_idx,
|
|
parameters,
|
|
g_group,
|
|
self.accumulate_steps,
|
|
act=HOOK_ACTION.REDUCE_SCATTER,
|
|
release_grads=self.sd_release_grads,
|
|
use_reduce_avg=self.use_reduce_avg,
|
|
free_grads_in_comm=False,
|
|
init_slice_param=False,
|
|
slice_params=self._slice_params,
|
|
)
|
|
group_idx += 1
|
|
self._comm_buffer_list.append(buffer)
|
|
if g_color not in self._color_to_comm_buffer_list.keys():
|
|
self._color_to_comm_buffer_list[g_color] = []
|
|
self._color_to_comm_buffer_list[g_color].append(buffer)
|
|
for p in parameters:
|
|
if p.name in self.param2bucket:
|
|
self.param2bucket[p.name].append(buffer)
|
|
else:
|
|
self.param2bucket[p.name] = [buffer]
|
|
|
|
self._comm_buffer_list.sort(key=lambda x: x._dst)
|
|
|
|
def clear_param_storage(self, color):
|
|
assert self._multi_precision, (
|
|
"Muon Sharding Optimizer only support clear param with multi_precision mode"
|
|
)
|
|
|
|
self.clear_color.add(color)
|
|
# 1D params
|
|
if color in self._color_to_comm_buffer_list.keys():
|
|
for comm_buffer in self._color_to_comm_buffer_list[color]:
|
|
has_clear = False
|
|
for param in comm_buffer.params:
|
|
grad_view = comm_buffer._sharding_param_grad_view[
|
|
param.name
|
|
]
|
|
slice_param = self._slice_params[param.name]
|
|
if (
|
|
not g_shard_bypass_dygraph_optimizer
|
|
and grad_view._param_begin < grad_view._param_end
|
|
):
|
|
grad_view.fill_slice_param(slice_param)
|
|
self._create_master_weight(slice_param)
|
|
if param.name in self._master_weights:
|
|
slice_param._clear_dataptr()
|
|
has_clear = True
|
|
|
|
if has_clear:
|
|
comm_buffer._clear_param_storage()
|
|
# 2D params
|
|
if color in self._params_2d_by_color.keys():
|
|
group_info = self._color_to_group_info[color]
|
|
sharding_rank = max(group_info["rank"], 0)
|
|
rank2params_2d_by_color = self._rank2params_2d_by_color[color]
|
|
local_2d = rank2params_2d_by_color[sharding_rank]
|
|
for param in local_2d:
|
|
if not g_shard_bypass_dygraph_optimizer:
|
|
self._create_master_weight(param)
|
|
|
|
for param in self._params_2d_by_color[color]:
|
|
param._clear_to_zero_allocation()
|
|
|
|
def reset_param_storage(self):
|
|
for color in self.clear_color:
|
|
if color is None:
|
|
continue
|
|
# 1D params
|
|
if color in self._color_to_comm_buffer_list.keys():
|
|
for comm_buffer in self._color_to_comm_buffer_list[color]:
|
|
if not comm_buffer.param_storage._is_initialized():
|
|
comm_buffer._reset_param_storage()
|
|
# 2D params
|
|
if color in self._params_2d_by_color.keys():
|
|
for param in self._params_2d_by_color[color]:
|
|
if not param._is_initialized():
|
|
new_param = paddle.empty_like(param)
|
|
new_param._share_buffer_to(param)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Gradient communication
|
|
# ------------------------------------------------------------------
|
|
|
|
def _get_param_grad(self, param):
|
|
if not param.trainable:
|
|
return None
|
|
if hasattr(param, "main_grad"):
|
|
assert param._grad_ivar() is None, (
|
|
"param.grad should be None when using main_grad"
|
|
)
|
|
return param.main_grad
|
|
return param._grad_ivar()
|
|
|
|
def _reduce_2d_grads(self, params, param2rank, comm_group):
|
|
"""Reduce gradients for 2D params to their owner rank within comm_group."""
|
|
for param in params:
|
|
g_var = self._get_param_grad(param)
|
|
if g_var is None:
|
|
if hasattr(param, "main_grad"):
|
|
g_var = paddle.zeros_like(param, dtype=paddle.float32)
|
|
param.main_grad = g_var
|
|
else:
|
|
g_var = paddle.zeros_like(param, dtype=param.dtype)
|
|
param.grad = g_var
|
|
|
|
reduce_op = ReduceOp.AVG
|
|
if not self.use_reduce_avg:
|
|
nranks = comm_group.nranks
|
|
g_var.scale_(1.0 / nranks)
|
|
reduce_op = ReduceOp.SUM
|
|
|
|
if paddle.distributed.in_auto_parallel_align_mode():
|
|
reduce_op = ReduceOp.SUM
|
|
|
|
param_rank = param2rank[param.name]
|
|
paddle.distributed.reduce(
|
|
g_var,
|
|
dst=comm_group.ranks[param_rank],
|
|
op=reduce_op,
|
|
group=comm_group,
|
|
sync_op=True,
|
|
)
|
|
|
|
def reduce_gradients(self, parameter_list, hcg):
|
|
"""Reduce gradients: reduce for 2D params, reduce-scatter for 1D params."""
|
|
if (
|
|
paddle.is_compiled_with_xpu()
|
|
and os.getenv("XPU_CDNN_CLUSTER_PARALLEL") is not None
|
|
):
|
|
paddle.device.synchronize()
|
|
|
|
with framework.no_grad():
|
|
# --- 2D params: reduce to owner rank via each color's group ---
|
|
if self._use_fuse_gradients:
|
|
for comm_buffer in self.comm_buffer_2d:
|
|
if (
|
|
self.sd_release_grads
|
|
and comm_buffer.grad_storage is None
|
|
):
|
|
if comm_buffer.need_reduce_scale_sync():
|
|
for param in comm_buffer.params:
|
|
comm_buffer._copy_grad_to_buffer(param)
|
|
if not self.use_group_call_opt:
|
|
for comm_buffer in self.comm_buffer_2d:
|
|
comm_buffer._comm_grads()
|
|
else:
|
|
same_card_buffers = []
|
|
all_ring_buffers = []
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
sharding_group = hcg.get_sharding_parallel_group()
|
|
moe_sharding_group = hcg.get_moe_sharding_parallel_group()
|
|
for comm_buffer in self.comm_buffer_2d:
|
|
if comm_buffer._comm_group == sharding_group:
|
|
all_ring_buffers.append(comm_buffer)
|
|
elif comm_buffer._comm_group == moe_sharding_group:
|
|
same_card_buffers.append(comm_buffer)
|
|
else:
|
|
raise ValueError("Unknown comm group")
|
|
tasks = []
|
|
with _coalescing_manager(moe_sharding_group, tasks):
|
|
for comm_buffer in all_ring_buffers:
|
|
dst_rank = get_same_card_rank(
|
|
list(moe_sharding_group.ranks), comm_buffer._dst
|
|
)
|
|
assert dst_rank != -1, "Please check you dst_rank!"
|
|
assert comm_buffer._use_reduce_avg, (
|
|
"Only support for reduce avg now"
|
|
)
|
|
paddle.distributed.stream.reduce(
|
|
comm_buffer.grad_storage,
|
|
dst=dst_rank,
|
|
op=paddle.distributed.ReduceOp.AVG,
|
|
group=moe_sharding_group,
|
|
sync_op=True,
|
|
use_calc_stream=True,
|
|
)
|
|
|
|
for comm_buffer in same_card_buffers:
|
|
assert (
|
|
comm_buffer._comm_group == moe_sharding_group
|
|
), "Please check comm group"
|
|
assert comm_buffer._use_reduce_avg, (
|
|
"Only support for reduce avg now"
|
|
)
|
|
paddle.distributed.stream.reduce(
|
|
comm_buffer.grad_storage,
|
|
dst=comm_buffer._dst,
|
|
op=paddle.distributed.ReduceOp.AVG,
|
|
group=comm_buffer._comm_group,
|
|
sync_op=True,
|
|
use_calc_stream=True,
|
|
)
|
|
|
|
rank = paddle.distributed.get_rank()
|
|
trainer_idx = rank // 8
|
|
trainer_comm_group = self.trainer_comms[trainer_idx]
|
|
with _coalescing_manager(trainer_comm_group, tasks):
|
|
for comm_buffer in all_ring_buffers:
|
|
trainer, trainer_ranks = get_trainer_ranks(
|
|
comm_buffer._dst
|
|
)
|
|
if rank in trainer_ranks:
|
|
comm_group = self.trainer_comms[trainer]
|
|
assert comm_group == trainer_comm_group, (
|
|
"Please check comm group"
|
|
)
|
|
assert comm_buffer._use_reduce_avg, (
|
|
"Only support for reduce avg now"
|
|
)
|
|
paddle.distributed.stream.reduce(
|
|
comm_buffer.grad_storage,
|
|
dst=comm_buffer._dst,
|
|
op=paddle.distributed.ReduceOp.AVG,
|
|
group=comm_group,
|
|
sync_op=True,
|
|
use_calc_stream=True,
|
|
)
|
|
|
|
else:
|
|
# --- 2D params: reduce to owner rank via each color's group ---
|
|
sharding_group = hcg.get_sharding_parallel_group()
|
|
for color_key, params_2d in self._params_2d_by_color.items():
|
|
if not params_2d:
|
|
continue
|
|
group_info = self._color_to_group_info.get(color_key, {})
|
|
group = group_info.get('group', sharding_group)
|
|
world_size = group_info.get('world_size', 1)
|
|
if world_size > 1:
|
|
param2rank = self._param2rank_2d_by_color.get(
|
|
color_key, {}
|
|
)
|
|
self._reduce_2d_grads(params_2d, param2rank, group)
|
|
|
|
# --- 1D params: reduce-scatter via comm buffers ---
|
|
for comm_buffer in self._comm_buffer_list:
|
|
if self.sd_release_grads and comm_buffer.grad_storage is None:
|
|
if comm_buffer.need_reduce_scale_sync():
|
|
for param in comm_buffer.params:
|
|
comm_buffer._copy_grad_to_buffer(param)
|
|
|
|
if not self.comm_overlap:
|
|
comm_buffer._comm_grads()
|
|
|
|
# wait for all comm_buffer tasks to finish
|
|
if self._use_fuse_gradients:
|
|
if not self.use_group_call_opt:
|
|
for comm_buffer in self.comm_buffer_2d:
|
|
comm_buffer.scale_grads()
|
|
|
|
for comm_buffer in self._comm_buffer_list:
|
|
comm_buffer.scale_grads()
|
|
|
|
def filter_parameters(self, parameter_list, hcg):
|
|
"""Filter parameters: return local 2D params + initialized 1D slices."""
|
|
local_1d = [
|
|
self._slice_params[p.name]
|
|
for p in parameter_list
|
|
if p.name in self._slice_params
|
|
]
|
|
local_1d = [p for p in local_1d if p._is_initialized()]
|
|
return self._local_2d + local_1d
|
|
|
|
# ------------------------------------------------------------------
|
|
# Parameter sync after optimizer step
|
|
# ------------------------------------------------------------------
|
|
def _broadcast_2d_params(self, rank2params, comm_group):
|
|
"""Broadcast 2D params from owner ranks within comm_group."""
|
|
broadcast_tasks = []
|
|
for rank, params in rank2params.items():
|
|
src_rank = comm_group.ranks[rank]
|
|
|
|
if len(params) == 0:
|
|
continue
|
|
|
|
with framework.no_grad():
|
|
# Calculate total size and build param-to-offset mapping
|
|
param_sizes = [np.prod(p.shape) for p in params]
|
|
total_size = sum(param_sizes)
|
|
dtype = params[0].dtype
|
|
|
|
param_buffer = paddle.empty(shape=[total_size], dtype=dtype)
|
|
offset = 0
|
|
for param in params:
|
|
param_shape = param.shape
|
|
stop_gradient = param.stop_gradient
|
|
param.stop_gradient = True
|
|
param.flatten_()
|
|
|
|
paddle.assign(
|
|
param,
|
|
param_buffer._slice(
|
|
offset, offset + np.prod(param_shape)
|
|
),
|
|
)
|
|
|
|
param.get_tensor()._set_dims(param_shape)
|
|
param.stop_gradient = stop_gradient
|
|
param_buffer._slice(
|
|
offset, offset + np.prod(param_shape)
|
|
)._share_buffer_to(param)
|
|
|
|
offset += np.prod(param_shape)
|
|
task = paddle.distributed.broadcast(
|
|
param_buffer,
|
|
src=src_rank,
|
|
group=comm_group,
|
|
sync_op=False,
|
|
)
|
|
broadcast_tasks.append(task)
|
|
return broadcast_tasks
|
|
|
|
def reorder_params(self, comm_list):
|
|
"""Reorder and flatten parameters into contiguous buffers for communication.
|
|
|
|
For each rank in the communication groups, flatten multiple parameters into a
|
|
single contiguous buffer to enable efficient communication operations
|
|
|
|
Args:
|
|
comm_list: List of tuples (rank2params, comm_group), where:
|
|
- rank2params: Dict mapping local rank indices to lists of parameters
|
|
- comm_group: Communication group containing ranks mapping
|
|
|
|
Returns:
|
|
List of dicts with keys:
|
|
- param_buffer: Contiguous buffer containing flattened parameters
|
|
- src_rank: Source rank for this buffer in the comm_group
|
|
"""
|
|
reorder_params_list = []
|
|
for rank2params, group in comm_list:
|
|
for rank, params in rank2params.items():
|
|
src_rank = group.ranks[rank]
|
|
if len(params) == 0:
|
|
continue
|
|
|
|
with framework.no_grad():
|
|
# Calculate total size and build param-to-offset mapping
|
|
param_sizes = [np.prod(p.shape) for p in params]
|
|
total_size = sum(param_sizes)
|
|
dtype = params[0].dtype
|
|
|
|
if len(params) != 1:
|
|
param_buffer = paddle.empty(
|
|
shape=[total_size], dtype=dtype
|
|
)
|
|
offset = 0
|
|
for param in params:
|
|
param_shape = param.shape
|
|
stop_gradient = param.stop_gradient
|
|
param.stop_gradient = True
|
|
param.flatten_()
|
|
|
|
paddle.assign(
|
|
param,
|
|
param_buffer._slice(
|
|
offset, offset + np.prod(param_shape)
|
|
),
|
|
)
|
|
|
|
param.get_tensor()._set_dims(param_shape)
|
|
param.stop_gradient = stop_gradient
|
|
param_buffer._slice(
|
|
offset, offset + np.prod(param_shape)
|
|
)._share_buffer_to(param)
|
|
|
|
offset += np.prod(param_shape)
|
|
else:
|
|
param_buffer = params[0]
|
|
reorder_params_item = {
|
|
"param_buffer": param_buffer,
|
|
"src_rank": src_rank,
|
|
}
|
|
reorder_params_list.append(reorder_params_item)
|
|
return reorder_params_list
|
|
|
|
def _sharding_sync_parameters(self):
|
|
"""Sync parameters: broadcast 2D, all-gather 1D."""
|
|
comm_group = self._hcg.get_sharding_parallel_group()
|
|
|
|
with framework.no_grad():
|
|
all_tasks = []
|
|
|
|
# --- 2D params: broadcast from owner via each color's group ---
|
|
all_ring_list = []
|
|
same_card_list = []
|
|
for color_key, rank2params in self._rank2params_2d_by_color.items():
|
|
group_info = self._color_to_group_info.get(color_key, {})
|
|
group = group_info.get('group', comm_group)
|
|
world_size = group_info.get('world_size', 1)
|
|
if world_size > 1:
|
|
if not self.use_group_call_opt:
|
|
all_tasks.extend(
|
|
self._broadcast_2d_params(rank2params, group)
|
|
)
|
|
else:
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
sharding_group = hcg.get_sharding_parallel_group()
|
|
moe_sharding_group = (
|
|
hcg.get_moe_sharding_parallel_group()
|
|
)
|
|
if group == sharding_group:
|
|
all_ring_list.append([rank2params, group])
|
|
elif group == moe_sharding_group:
|
|
same_card_list.append([rank2params, group])
|
|
else:
|
|
raise ValueError("Please check your comm group")
|
|
# world_size=1: single rank group, no broadcast needed
|
|
|
|
if self.use_group_call_opt:
|
|
all_ring_reorder_params_list = self.reorder_params(
|
|
all_ring_list
|
|
)
|
|
same_card_reorder_params_list = self.reorder_params(
|
|
same_card_list
|
|
)
|
|
rank = paddle.distributed.get_rank()
|
|
|
|
# trainer inner broadcast
|
|
rank = paddle.distributed.get_rank()
|
|
trainer_idx = rank // 8
|
|
trainer_comm_group = self.trainer_comms[trainer_idx]
|
|
tasks = []
|
|
with _coalescing_manager(trainer_comm_group, tasks):
|
|
for reorder_params_item in all_ring_reorder_params_list:
|
|
param_buffer = reorder_params_item["param_buffer"]
|
|
src_rank = reorder_params_item["src_rank"]
|
|
trainer, trainer_ranks = get_trainer_ranks(src_rank)
|
|
if rank in trainer_ranks:
|
|
comm_group = self.trainer_comms[trainer]
|
|
assert comm_group == trainer_comm_group, (
|
|
"Please check comm group"
|
|
)
|
|
paddle.distributed.stream.broadcast(
|
|
param_buffer,
|
|
src=src_rank,
|
|
group=comm_group,
|
|
sync_op=True,
|
|
use_calc_stream=True,
|
|
)
|
|
|
|
# same card broadcast
|
|
with _coalescing_manager(moe_sharding_group, tasks):
|
|
for reorder_params_item in all_ring_reorder_params_list:
|
|
param_buffer = reorder_params_item["param_buffer"]
|
|
src_rank = reorder_params_item["src_rank"]
|
|
same_card_src_rank = get_same_card_rank(
|
|
list(moe_sharding_group.ranks), src_rank
|
|
)
|
|
assert same_card_src_rank != -1, (
|
|
"Please check your same_card_src_rank in broadcast!"
|
|
)
|
|
paddle.distributed.stream.broadcast(
|
|
param_buffer,
|
|
src=same_card_src_rank,
|
|
group=moe_sharding_group,
|
|
sync_op=True,
|
|
use_calc_stream=True,
|
|
)
|
|
for same_card_item in same_card_reorder_params_list:
|
|
param_buffer = same_card_item["param_buffer"]
|
|
src_rank = same_card_item["src_rank"]
|
|
paddle.distributed.stream.broadcast(
|
|
param_buffer,
|
|
src=src_rank,
|
|
group=moe_sharding_group,
|
|
sync_op=True,
|
|
use_calc_stream=True,
|
|
)
|
|
|
|
if not self.use_group_call_opt:
|
|
for task in all_tasks:
|
|
task.wait()
|
|
|
|
# --- 1D params: all-gather via comm buffers ---
|
|
for comm_buffer in self._comm_buffer_list:
|
|
comm_buffer.sync_params()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Clear gradients
|
|
# ------------------------------------------------------------------
|
|
|
|
def clear_grad(self, set_to_zero=True):
|
|
"""Clear gradients for all parameters."""
|
|
|
|
def clear_grad_func(p):
|
|
if hasattr(p, "main_grad") and p.main_grad is not None:
|
|
assert p._grad_ivar() is None
|
|
if set_to_zero:
|
|
p.main_grad.zero_()
|
|
else:
|
|
p.main_grad._clear()
|
|
p.main_grad = None
|
|
elif not hasattr(p, "main_grad"):
|
|
if self.tensor_fusion:
|
|
if set_to_zero:
|
|
p.grad.zero_()
|
|
else:
|
|
p.grad._clear()
|
|
p.grad = None
|
|
else:
|
|
p.clear_gradient(set_to_zero)
|
|
|
|
for p in self._parameter_list:
|
|
clear_grad_func(p)
|
|
|
|
if self.sd_release_grads and not self.pp_overlap:
|
|
# 1D params are managed by comm buffers
|
|
for comm_buffer in self._comm_buffer_list:
|
|
if comm_buffer.need_reduce_scale_sync():
|
|
comm_buffer._clear_grad_storage()
|
|
# 2D params are managed by comm buffers
|
|
if self._use_fuse_gradients:
|
|
for comm_buffer in self.comm_buffer_2d:
|
|
if comm_buffer.need_reduce_scale_sync():
|
|
comm_buffer._clear_grad_storage()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Optimizer step
|
|
# ------------------------------------------------------------------
|
|
|
|
def _collect_comm_buffers(self):
|
|
"""Collect communication buffers (for PP overlap compatibility)."""
|
|
if self._comm_buffer_list:
|
|
return
|
|
for param in self._params_1d:
|
|
if not hasattr(param, "comm_buffer_ref"):
|
|
continue
|
|
comm_buffer_ref = param.comm_buffer_ref
|
|
del param.comm_buffer_ref
|
|
comm_buffer = comm_buffer_ref()
|
|
self._comm_buffer_list.append(comm_buffer)
|
|
|
|
for bucket in self._comm_buffer_list:
|
|
for p in bucket._params:
|
|
if p.name in self.param2bucket:
|
|
self.param2bucket[p.name].append(bucket)
|
|
else:
|
|
self.param2bucket[p.name] = [bucket]
|
|
|
|
def _assign_slice_grad(self):
|
|
"""Assign gradients from comm buffers to slice params for 1D params."""
|
|
for comm_buffer in self._comm_buffer_list:
|
|
for param in comm_buffer.params:
|
|
if param.name in self._slice_params:
|
|
slice_param = self._slice_params[param.name]
|
|
if self.sd_release_grads and hasattr(
|
|
slice_param, "main_grad"
|
|
):
|
|
if not slice_param.main_grad._is_initialized():
|
|
del slice_param.main_grad
|
|
comm_buffer.assign_slice_grad(param, slice_param)
|
|
|
|
def set_grad_clip_info(self):
|
|
from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.hybrid_parallel_optimizer import (
|
|
ClipGradByAdaptiveNorm,
|
|
)
|
|
|
|
if isinstance(self._inner_opt._grad_clip, ClipGradByAdaptiveNorm):
|
|
if not hasattr(
|
|
self._inner_opt._grad_clip, "group_index_to_name_to_local_idx"
|
|
):
|
|
# Build parameter groups and their metadata
|
|
group_index_to_name_to_local_idx = {} # group_idx -> {param_name: local_index}
|
|
param_name_to_group_idx = {} # param_name -> group_idx
|
|
group_idx_to_comm_group = {} # group_idx -> comm_group
|
|
group_idx_to_param_num = {} # group_idx -> param count
|
|
|
|
group_idx = 0
|
|
# 1D params
|
|
for _, comm_buffer in enumerate(self._comm_buffer_list):
|
|
name_to_local_idx = {}
|
|
for local_idx, param in enumerate(comm_buffer._params):
|
|
param = self._slice_params[param.name]
|
|
if param._is_initialized():
|
|
name_to_local_idx[param.name] = local_idx
|
|
param_name_to_group_idx[param.name] = group_idx
|
|
group_index_to_name_to_local_idx[group_idx] = (
|
|
name_to_local_idx
|
|
)
|
|
group_idx_to_comm_group[group_idx] = comm_buffer._comm_group
|
|
group_idx_to_param_num[group_idx] = len(comm_buffer._params)
|
|
group_idx += 1
|
|
|
|
# 2D params
|
|
for color_key, params_2d in self._params_2d_by_color.items():
|
|
name_to_local_idx = {}
|
|
local_2d_names = {p.name for p in self._local_2d}
|
|
for local_idx, param in enumerate(params_2d):
|
|
if param.name in local_2d_names:
|
|
name_to_local_idx[param.name] = local_idx
|
|
param_name_to_group_idx[param.name] = group_idx
|
|
group_index_to_name_to_local_idx[group_idx] = (
|
|
name_to_local_idx
|
|
)
|
|
|
|
# Get comm group from color_key
|
|
group_info = self._color_to_group_info[color_key]
|
|
comm_group = group_info['group']
|
|
|
|
group_idx_to_comm_group[group_idx] = comm_group
|
|
group_idx_to_param_num[group_idx] = len(params_2d)
|
|
group_idx += 1
|
|
|
|
self._inner_opt._grad_clip.group_index_to_name_to_local_idx = (
|
|
group_index_to_name_to_local_idx
|
|
)
|
|
self._inner_opt._grad_clip.param_name_to_group_idx = (
|
|
param_name_to_group_idx
|
|
)
|
|
self._inner_opt._grad_clip.group_idx_to_comm_group = (
|
|
group_idx_to_comm_group
|
|
)
|
|
self._inner_opt._grad_clip.group_idx_to_param_num = (
|
|
group_idx_to_param_num
|
|
)
|
|
self._inner_opt._grad_clip.sharding_stage1_v2 = True
|
|
|
|
def step(self):
|
|
"""Optimizer step: update local 2D params and 1D slices, then sync."""
|
|
self.reset_param_storage()
|
|
|
|
self._collect_comm_buffers()
|
|
self._assign_slice_grad()
|
|
|
|
self.set_grad_clip_info()
|
|
|
|
if not isinstance(self._origin_parameter_list[0], dict):
|
|
params_grads = []
|
|
|
|
# --- All 2D params on this rank (all colors): full tensors ---
|
|
# Pass the original param directly to the optimizer.
|
|
# _muon_update handles both shapes:
|
|
# - 2D [H, I]: standard Newton-Schulz
|
|
# - 3D [n_experts, H, I]: per-expert Newton-Schulz loop
|
|
# Keeping the original param name avoids registering _expert_N
|
|
# accumulator keys absent from model_sharded_state_dict, which
|
|
# would break sharded_state_dict (checkpoint save).
|
|
for color_key, rank2params in self._rank2params_2d_by_color.items():
|
|
group_info = self._color_to_group_info.get(color_key, {})
|
|
color_rank = group_info.get('rank', 0)
|
|
world_size = group_info.get('world_size', 1)
|
|
rank_key = color_rank if world_size > 1 else 0
|
|
for param in rank2params.get(rank_key, []):
|
|
if param.stop_gradient:
|
|
continue
|
|
grad_var = param._grad_ivar()
|
|
if (
|
|
hasattr(param, "main_grad")
|
|
and param.main_grad is not None
|
|
):
|
|
grad_var = param.main_grad
|
|
if grad_var is not None:
|
|
params_grads.append((param, grad_var))
|
|
|
|
# --- 1D params: slice params (element-wise shards) ---
|
|
for param in self._params_1d:
|
|
if param.stop_gradient:
|
|
continue
|
|
if param.name not in self._slice_params:
|
|
continue
|
|
slice_p = self._slice_params[param.name]
|
|
if not slice_p._is_initialized():
|
|
continue
|
|
grad_var = slice_p._grad_ivar()
|
|
if (
|
|
hasattr(slice_p, "main_grad")
|
|
and slice_p.main_grad is not None
|
|
):
|
|
grad_var = slice_p.main_grad
|
|
if grad_var is not None:
|
|
params_grads.append((slice_p, grad_var))
|
|
|
|
self._apply_optimize(
|
|
loss=None,
|
|
startup_program=None,
|
|
params_grads=params_grads,
|
|
)
|
|
|
|
# Sync parameters across sharding ranks
|
|
self._sharding_sync_parameters()
|
|
|
|
# ------------------------------------------------------------------
|
|
# State dict (checkpoint save/load)
|
|
# ------------------------------------------------------------------
|
|
|
|
@framework.dygraph_only
|
|
def set_state_dict(self, state_dict):
|
|
inner_state = {}
|
|
# Collect local parameters: 2D whole-tensor params + 1D original params
|
|
parameters = list(self._local_2d) + list(self._params_1d)
|
|
|
|
if "LR_Scheduler" in state_dict:
|
|
inner_state["LR_Scheduler"] = state_dict.pop("LR_Scheduler")
|
|
|
|
if "master_weights" in state_dict:
|
|
master = state_dict.pop("master_weights")
|
|
inner_state["master_weights"] = {}
|
|
for p in parameters:
|
|
for k, v in master.items():
|
|
if p.name == k:
|
|
v.name = self._inner_opt._gen_master_weight_var_name(p)
|
|
inner_state["master_weights"][k] = v
|
|
|
|
for p in parameters:
|
|
for k, v in state_dict.items():
|
|
if p.name in k:
|
|
inner_state[k] = v
|
|
|
|
self._inner_opt.set_state_dict(inner_state)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Utility
|
|
# ------------------------------------------------------------------
|
|
|
|
def _set_inner_opt_attr(self, attr_name, value):
|
|
inner_opt = self._inner_opt
|
|
inner_opt_name = '_inner_opt'
|
|
if not isinstance(attr_name, str):
|
|
raise TypeError(
|
|
f"attr_name should be str type, but is {type(attr_name)}"
|
|
)
|
|
while hasattr(inner_opt, attr_name):
|
|
setattr(inner_opt, attr_name, value)
|
|
inner_opt = getattr(inner_opt, inner_opt_name, None)
|
|
if inner_opt is None:
|
|
break
|
|
|
|
def sharded_state_dict(self, model_sharded_state_dict):
|
|
"""Build a sharded optimizer state dict for flex checkpoint save/load.
|
|
|
|
Overrides the inner Muon optimizer's sharded_state_dict to handle V3's
|
|
hybrid sharding scheme:
|
|
- 2D Muon params: whole tensor, shape matches model's local_shape.
|
|
Handled by delegating to the inner Muon's sharded_state_dict after
|
|
filtering out 1D param states.
|
|
- 1D AdamW params: accumulators are 1D shards (from reduce-scatter);
|
|
wrapped with is_flattened=True + flattened_range, like V2.
|
|
"""
|
|
import paddle as _paddle
|
|
from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
|
|
ShardedWeight,
|
|
create_sharded_weight_with_new_local,
|
|
)
|
|
|
|
# ---- Step 1: Collect flattened_range for each 1D (AdamW) param ----
|
|
# Identical logic to DygraphShardingOptimizerV2.sharded_state_dict.
|
|
param_slice_info = {} # param_name -> slice(begin, end)
|
|
padded_param = set()
|
|
for buffer in self._comm_buffer_list:
|
|
for (
|
|
param_name,
|
|
grad_view,
|
|
) in buffer._sharding_param_grad_view.items():
|
|
numel = grad_view._param.numel().item()
|
|
param_begin = grad_view._param_begin
|
|
param_end = grad_view._param_end
|
|
index = grad_view._index
|
|
padding_begin = index + numel
|
|
flattened_range = slice(
|
|
param_begin - index,
|
|
max(
|
|
min(padding_begin - index, param_end - index),
|
|
param_begin - index,
|
|
),
|
|
)
|
|
if param_end > padding_begin:
|
|
padded_param.add(param_name)
|
|
param_slice_info[param_name] = flattened_range
|
|
|
|
# ---- Step 2: Build static_name → struct_name mapping ----
|
|
model_sharded_sorted = dict(sorted(model_sharded_state_dict.items()))
|
|
static_to_struct = {}
|
|
for struct_name, sw in model_sharded_sorted.items():
|
|
if sw.local_tensor.name not in static_to_struct:
|
|
static_to_struct[sw.local_tensor.name] = struct_name
|
|
|
|
# ---- Step 3: Process all optimizer states ----
|
|
_FP32_MASTER = "fp32_master_0"
|
|
_optimizer_scalar_names = ["beta1_pow_acc_0", "beta2_pow_acc_0"]
|
|
_optimizer_vector_names = ["moment1_0", "moment2_0"]
|
|
|
|
def _make_2d_entry(uname, t, sp):
|
|
"""Reshape tensor if numel matches but shape differs, then wrap as ShardedWeight."""
|
|
target = sp.local_shape
|
|
if (
|
|
tuple(t.shape) != tuple(target)
|
|
and t.numel() == _paddle.to_tensor(list(target)).prod().item()
|
|
):
|
|
t = t.reshape(target)
|
|
return create_sharded_weight_with_new_local(uname, t, sp)
|
|
|
|
def _split_state_name(vname):
|
|
if _FP32_MASTER in vname:
|
|
return tuple(vname.split("_" + _FP32_MASTER + "_", 1))
|
|
for suffix in _optimizer_scalar_names + _optimizer_vector_names:
|
|
if vname.endswith(suffix):
|
|
return vname[: -(len(suffix) + 1)], suffix
|
|
raise ValueError(
|
|
f"Cannot parse optimizer state variable name: {vname!r}"
|
|
)
|
|
|
|
optimizer_state_dict = self._inner_opt.state_dict()
|
|
master_weights = optimizer_state_dict.pop("master_weights", None)
|
|
optimizer_state_dict.pop("LR_Scheduler", None)
|
|
|
|
sharded_state = {}
|
|
|
|
for key, tensor in optimizer_state_dict.items():
|
|
static_name, state_type = _split_state_name(key)
|
|
if static_name not in static_to_struct:
|
|
continue
|
|
|
|
struct_name = static_to_struct[static_name]
|
|
sharded_param = model_sharded_sorted[struct_name]
|
|
unified_name = f"{struct_name}.{state_type}"
|
|
|
|
is_1d_param = static_name in param_slice_info
|
|
|
|
if state_type in _optimizer_vector_names:
|
|
if is_1d_param:
|
|
# 1D AdamW shard: wrap with is_flattened=True (like V2)
|
|
flattened_range = param_slice_info[static_name]
|
|
if flattened_range.stop - flattened_range.start == 0:
|
|
continue
|
|
is_padded = static_name in padded_param
|
|
if is_padded:
|
|
local_tensor = _paddle.slice(
|
|
tensor,
|
|
axes=[0],
|
|
starts=[0],
|
|
ends=[flattened_range.stop - flattened_range.start],
|
|
)
|
|
else:
|
|
local_tensor = tensor
|
|
sharded_state[unified_name] = ShardedWeight(
|
|
key=unified_name,
|
|
local_tensor=local_tensor,
|
|
local_shape=sharded_param.local_shape,
|
|
global_shape=sharded_param.global_shape,
|
|
global_offset=sharded_param.global_offset,
|
|
is_flattened=True,
|
|
flattened_range=flattened_range,
|
|
)
|
|
elif tensor.is_dist():
|
|
sharded_state[unified_name] = ShardedWeight(
|
|
key=unified_name,
|
|
local_tensor=tensor,
|
|
local_shape=tensor.shape,
|
|
global_shape=tensor.shape,
|
|
global_offset=sharded_param.global_offset,
|
|
)
|
|
else:
|
|
# 2D Muon param (non-MoE or MoE): shape may differ between
|
|
# Python param.shape (3D view) and model storage (2D).
|
|
sharded_state[unified_name] = _make_2d_entry(
|
|
unified_name, tensor, sharded_param
|
|
)
|
|
else:
|
|
# Scalar states (beta_pow): replicated
|
|
sharded_state[unified_name] = ShardedWeight(
|
|
key=unified_name,
|
|
local_tensor=tensor,
|
|
local_shape=(1,),
|
|
global_shape=(1,),
|
|
global_offset=(0,),
|
|
)
|
|
|
|
# FP32 master weights
|
|
if master_weights:
|
|
for weight_key, tensor in master_weights.items():
|
|
if weight_key not in static_to_struct:
|
|
continue
|
|
struct_name = static_to_struct[weight_key]
|
|
sharded_param = model_sharded_sorted[struct_name]
|
|
unified_name = f"{struct_name}.w_0"
|
|
is_1d_param = weight_key in param_slice_info
|
|
|
|
if is_1d_param:
|
|
flattened_range = param_slice_info[weight_key]
|
|
if flattened_range.stop - flattened_range.start == 0:
|
|
continue
|
|
is_padded = weight_key in padded_param
|
|
if is_padded:
|
|
local_tensor = _paddle.slice(
|
|
tensor,
|
|
axes=[0],
|
|
starts=[0],
|
|
ends=[flattened_range.stop - flattened_range.start],
|
|
)
|
|
else:
|
|
local_tensor = tensor
|
|
sharded_state[unified_name] = ShardedWeight(
|
|
key=unified_name,
|
|
local_tensor=local_tensor,
|
|
local_shape=sharded_param.local_shape,
|
|
global_shape=sharded_param.global_shape,
|
|
global_offset=sharded_param.global_offset,
|
|
is_flattened=True,
|
|
flattened_range=flattened_range,
|
|
)
|
|
elif tensor.is_dist():
|
|
sharded_state[unified_name] = ShardedWeight(
|
|
key=unified_name,
|
|
local_tensor=tensor,
|
|
local_shape=tensor.shape,
|
|
global_shape=tensor.shape,
|
|
global_offset=sharded_param.global_offset,
|
|
)
|
|
else:
|
|
# Same reshape logic as for optimizer vector states:
|
|
# FP32 master weight may be 3D (e.g. grouped_gemm_experts
|
|
# [n_experts, H, I]) while model storage is 2D [n_experts*H, I].
|
|
sharded_state[unified_name] = _make_2d_entry(
|
|
unified_name, tensor, sharded_param
|
|
)
|
|
|
|
return sharded_state
|
|
|
|
def __getattr__(self, item):
|
|
return getattr(self._inner_opt, item)
|