initial clean commit

1faccd4 6 months ago

12.8 kB

# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# 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.
import importlib
import logging
import re
from collections import defaultdict
from collections.abc import Generator, Iterator
from dataclasses import dataclass
from typing import Any

import torch
import torch.distributed
import torch.nn as nn
from torch.distributed._composable.fsdp import FSDPModule
from torch.distributed.tensor import DTensor
from torch.nn.attention.flex_attention import _mask_mod_signature, and_masks
from torchtitan.components.dataloader import BaseDataLoader
from torchtitan.models.common.attention import (
    AttentionMasksType,
    VarlenMetadata,
    create_attention_mask,
    get_causal_mask_mod,
)

logger = logging.getLogger(__name__)


class NoOpDataLoader(BaseDataLoader):
    """A no-op dataloader for use when verl manages its own data loading.
    Satisfies the BaseDataLoader interface required by torchtitan's Trainer
    but does nothing. Its __iter__ yields nothing, and state_dict /
    load_state_dict are no-ops.
    """

    @dataclass(kw_only=True, slots=True)
    class Config(BaseDataLoader.Config):
        pass

    def __init__(self, **kwargs):
        pass

    def __iter__(self) -> Iterator[tuple[dict[str, torch.Tensor], torch.Tensor]]:
        return iter([])

    def state_dict(self):
        return {}

    def load_state_dict(self, state_dict):
        pass


# Mapping from HuggingFace model_type to torchtitan model name.
# Torchtitan models not mapped here:
#   - flux: diffusion model, not applicable to verl's RL/SFT workflows
#   - llama3_ft: fault-tolerant variant of llama3, same HF models (mapped via "llama")
_HF_MODEL_TYPE_TO_TORCHTITAN_NAME = {
    "qwen2": "qwen3",
    "qwen3": "qwen3",
    "qwen2_moe": "qwen3",
    "qwen3_moe": "qwen3",
    "llama": "llama3",
    "llama4": "llama4",
    "deepseek_v3": "deepseek_v3",
    "gpt_oss": "gpt_oss",
}


def derive_torchtitan_name_and_flavor(hf_config) -> tuple[str, str]:
    """Derive torchtitan model name and flavor from a HuggingFace config.
    The name is mapped from ``hf_config.model_type``. The flavor is found by
    matching architecture parameters (dim, n_layers, vocab_size) against the
    known flavors registered in the torchtitan model package.
    Args:
        hf_config: A HuggingFace AutoConfig object.
    Returns:
        A ``(name, flavor)`` tuple.
    Raises:
        ValueError: If model_type is unsupported or no matching flavor is found.
    """
    model_type = getattr(hf_config, "model_type", None)
    if model_type is None:
        raise ValueError("HuggingFace config does not have 'model_type' field")

    name = _HF_MODEL_TYPE_TO_TORCHTITAN_NAME.get(model_type)
    if name is None:
        raise ValueError(
            f"Cannot derive torchtitan model name from HF model_type '{model_type}'. "
            f"Supported types: {list(_HF_MODEL_TYPE_TO_TORCHTITAN_NAME.keys())}."
        )

    # Import the model package and find the configs dict
    model_module = importlib.import_module(f"torchtitan.models.{name}")
    model_configs = None
    for attr in dir(model_module):
        obj = getattr(model_module, attr)
        if isinstance(obj, dict) and attr.endswith("_configs"):
            model_configs = obj
            break

    if model_configs is None:
        raise ValueError(
            f"Could not find model configs dict in torchtitan.models.{name}. "
            f"Expected a dict attribute ending with '_configs'."
        )

    hidden_size = hf_config.hidden_size
    num_layers = hf_config.num_hidden_layers
    vocab_size = hf_config.vocab_size

    for flavor_name, model_cfg in model_configs.items():
        if (
            getattr(model_cfg, "dim", None) == hidden_size
            and getattr(model_cfg, "n_layers", None) == num_layers
            and getattr(model_cfg, "vocab_size", None) == vocab_size
        ):
            logger.info(
                f"Auto-derived torchtitan name='{name}', flavor='{flavor_name}' from HF model_type='{model_type}'"
            )
            return name, flavor_name

    raise ValueError(
        f"No matching torchtitan flavor found for model_type='{model_type}' "
        f"(hidden_size={hidden_size}, num_hidden_layers={num_layers}, "
        f"vocab_size={vocab_size}). "
        f"Available flavors for '{name}': {list(model_configs.keys())}."
    )


def enable_fsdp_gradient_division(model: nn.Module, dp_size: int) -> None:
    """
    Re-enable FSDP's automatic gradient division.
    TorchTitan calls disable_fsdp_gradient_division() which sets gradient_divide_factor=1.0.
    This re-enables it by setting the factor to the specified dp_size, so gradients are
    averaged across FSDP ranks. This is needed for verl's loss scaling (loss * dp_size)
    to work correctly.
    Args:
        model: The model (or model part) to enable gradient division on.
        dp_size: The data parallel size to use as the gradient divide factor.
    """

    for module in model.modules():
        if isinstance(module, FSDPModule):
            module.set_gradient_divide_factor(float(dp_size))


def get_attention_masks(
    input_batch: torch.Tensor,
    positions: torch.Tensor,
    attn_type: str,
) -> AttentionMasksType:
    match attn_type:
        case "flex":
            return _get_flex_attention_masks(
                input_batch,
                positions,
            )
        case "varlen":
            return _create_varlen_metadata_for_document(
                input_batch,
                positions,
            )
        case _:
            raise TypeError("Only varlen and flex attn masks are supported")


def _get_document_mask_mod(positions: torch.Tensor) -> _mask_mod_signature:
    # Detect boundaries from position resets
    first_dummy_value = positions[:, :1] - 1
    position_diff = torch.diff(positions, prepend=first_dummy_value, dim=-1)
    sequence_indices = (position_diff != 1).cumsum(-1)  # [batch, seq]

    def document_mask(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor) -> torch.Tensor:
        return sequence_indices[b, q_idx] == sequence_indices[b, kv_idx]

    return document_mask


def _get_flex_attention_masks(
    input_batch: torch.Tensor,
    positions: torch.Tensor,
) -> AttentionMasksType:
    mask_mods = [get_causal_mask_mod()]
    B = input_batch.shape[0]
    mask_mods.append(_get_document_mask_mod(positions=positions))
    return create_attention_mask(and_masks(*mask_mods), B, None, input_batch.shape[1], input_batch.shape[1])


def _create_varlen_metadata_for_document(input_batch: torch.Tensor, positions: torch.Tensor) -> VarlenMetadata:
    """
    Creates cumulative sequence length indices needed for variable length attention
    Args:
        input_batch: Input token IDs with shape [batch, seq].
        positions: Position IDs with shape [batch, seq]. Boundaries detected where
            position diff != 1 (i.e., position resets).
    Returns:
        VarlenMetadata containing cumulative sequence length indices for q, k, and max_seq_len
    """
    batch_size, seq_len = input_batch.shape
    device = input_batch.device

    # Detect boundaries from position resets (where diff != 1)
    first_dummy_value = positions[:, :1] - 1
    position_diff = torch.diff(positions, prepend=first_dummy_value, dim=-1)
    # boundary_mask[b, i] is True if position i starts a new document
    boundary_mask = position_diff != 1  # [batch, seq]
    boundary_mask[:, 0] = True

    cu_seqlens_list, all_seq_lengths = [], []
    offset = 0

    for b in range(batch_size):
        # Find positions where new documents start
        boundary_positions = boundary_mask[b].nonzero(as_tuple=True)[0].to(torch.int32)
        sample_cu_seqlens = torch.cat(
            [
                boundary_positions,
                torch.tensor([seq_len], dtype=torch.int32, device=device),
            ]
        )
        sample_cu_seqlens = torch.unique_consecutive(sample_cu_seqlens)

        seq_lengths = torch.diff(sample_cu_seqlens)
        all_seq_lengths.append(seq_lengths)

        cu_seqlens_adjusted = sample_cu_seqlens[:-1] + offset
        cu_seqlens_list.append(cu_seqlens_adjusted)

        offset += seq_len

    packed_cu_seqlens = torch.cat(cu_seqlens_list + [torch.tensor([offset], dtype=torch.int32, device=device)])

    max_seqlen = 0
    if len(all_seq_lengths) > 0:
        all_seq_lengths = torch.cat(all_seq_lengths)
        # device to host sync but only done once per model forward
        max_seqlen = all_seq_lengths.max().item()

    return VarlenMetadata(
        cu_seq_q=packed_cu_seqlens,
        cu_seq_k=packed_cu_seqlens,
        max_q=max_seqlen,
        max_k=max_seqlen,
    )


# Regex to parse: model.layers.{L}.mlp.experts.{E}.{weight_suffix}
_EXPERT_PATTERN = re.compile(r"\.layers\.(\d+)\..*\.experts\.(\d+)\.(.*)")


def _parse_expert_name(name: str) -> tuple[int, int, str] | None:
    """Parse layer_id, expert_id, weight_suffix from expert param name."""
    match = _EXPERT_PATTERN.search(name)
    if match:
        return int(match.group(1)), int(match.group(2)), match.group(3)
    return None


def _make_expert_name_template(name: str) -> str:
    """Convert 'model.layers.0.mlp.experts.3.w1' -> 'model.layers.0.mlp.experts.{}.w1'"""
    return _EXPERT_PATTERN.sub(lambda m: f".layers.{m.group(1)}.mlp.experts.{{}}.{m.group(3)}", name)


def iter_per_tensor_params_ep(
    params: dict[str, Any],
    device: int,
    ep_group: torch.distributed.ProcessGroup,
    ep_size: int,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Yield (name, tensor) pairs for weight sync with Expert Parallel.
    Gathers expert weights across EP ranks one (layer, weight_type) group
    at a time to avoid OOM from materializing all experts simultaneously.
    Non-expert params are yielded first (with FSDP full_tensor() if needed),
    then expert params are all-gathered per group and yielded individually.
    Args:
        params: HF-format state dict with per-expert keys. Expert keys must
            follow the pattern ``model.layers.{L}.mlp.experts.{E}.{suffix}``.
        device: device ID to place tensors on.
        ep_group: The EP process group for all-gather.
        ep_size: Number of EP ranks.
    """
    expert_params: dict[tuple[int, str], dict[int, tuple[str, Any]]] = defaultdict(dict)
    non_expert_params: list[tuple[str, Any]] = []

    for name, param in params.items():
        parsed = _parse_expert_name(name) if "mlp.experts." in name else None
        if parsed is None:
            non_expert_params.append((name, param))
        else:
            layer_id, expert_id, weight_suffix = parsed
            expert_params[(layer_id, weight_suffix)][expert_id] = (name, param)

    params.clear()

    # Yield non-expert params
    for name, param in non_expert_params:
        if isinstance(param, DTensor):
            yield name, param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True)
        else:
            yield name, param
    del non_expert_params

    # Yield expert params with all-gather
    for (layer_id, weight_suffix), experts_dict in sorted(expert_params.items()):
        sorted_expert_ids = sorted(experts_dict.keys())

        # Stack local expert weights
        local_weights = []
        for eid in sorted_expert_ids:
            _, param = experts_dict[eid]
            if isinstance(param, DTensor):
                param = param.to(device, non_blocking=True).full_tensor()
            else:
                param = param.to(device, non_blocking=True)
            local_weights.append(param)

        name_template = _make_expert_name_template(experts_dict[sorted_expert_ids[0]][0])
        local_stacked = torch.stack(local_weights, dim=0)

        # All-gather across EP ranks
        gathered_list = [torch.empty_like(local_stacked) for _ in range(ep_size)]
        torch.distributed.all_gather(gathered_list, local_stacked, group=ep_group)
        all_experts = torch.cat(gathered_list, dim=0)

        for expert_id in range(all_experts.shape[0]):
            yield name_template.format(expert_id), all_experts[expert_id].to(torch.bfloat16).clone()

        del local_weights, local_stacked, gathered_list, all_experts
        torch.cuda.empty_cache()

	# Copyright 2024 Bytedance Ltd. and/or its affiliates
	#
	# 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.
	import importlib
	import logging
	import re
	from collections import defaultdict
	from collections.abc import Generator, Iterator
	from dataclasses import dataclass
	from typing import Any

	import torch
	import torch.distributed
	import torch.nn as nn
	from torch.distributed._composable.fsdp import FSDPModule
	from torch.distributed.tensor import DTensor
	from torch.nn.attention.flex_attention import _mask_mod_signature, and_masks
	from torchtitan.components.dataloader import BaseDataLoader
	from torchtitan.models.common.attention import (
	AttentionMasksType,
	VarlenMetadata,
	create_attention_mask,
	get_causal_mask_mod,
	)

	logger = logging.getLogger(__name__)


	class NoOpDataLoader(BaseDataLoader):
	"""A no-op dataloader for use when verl manages its own data loading.

	Satisfies the BaseDataLoader interface required by torchtitan's Trainer
	but does nothing. Its __iter__ yields nothing, and state_dict /
	load_state_dict are no-ops.
	"""

	@dataclass(kw_only=True, slots=True)
	class Config(BaseDataLoader.Config):
	pass

	def __init__(self, **kwargs):
	pass

	def __iter__(self) -> Iterator[tuple[dict[str, torch.Tensor], torch.Tensor]]:
	return iter([])

	def state_dict(self):
	return {}

	def load_state_dict(self, state_dict):
	pass


	# Mapping from HuggingFace model_type to torchtitan model name.
	# Torchtitan models not mapped here:
	# - flux: diffusion model, not applicable to verl's RL/SFT workflows
	# - llama3_ft: fault-tolerant variant of llama3, same HF models (mapped via "llama")
	_HF_MODEL_TYPE_TO_TORCHTITAN_NAME = {
	"qwen2": "qwen3",
	"qwen3": "qwen3",
	"qwen2_moe": "qwen3",
	"qwen3_moe": "qwen3",
	"llama": "llama3",
	"llama4": "llama4",
	"deepseek_v3": "deepseek_v3",
	"gpt_oss": "gpt_oss",
	}


	def derive_torchtitan_name_and_flavor(hf_config) -> tuple[str, str]:
	"""Derive torchtitan model name and flavor from a HuggingFace config.

	The name is mapped from ``hf_config.model_type``. The flavor is found by
	matching architecture parameters (dim, n_layers, vocab_size) against the
	known flavors registered in the torchtitan model package.

	Args:
	hf_config: A HuggingFace AutoConfig object.

	Returns:
	A ``(name, flavor)`` tuple.

	Raises:
	ValueError: If model_type is unsupported or no matching flavor is found.
	"""
	model_type = getattr(hf_config, "model_type", None)
	if model_type is None:
	raise ValueError("HuggingFace config does not have 'model_type' field")

	name = _HF_MODEL_TYPE_TO_TORCHTITAN_NAME.get(model_type)
	if name is None:
	raise ValueError(
	f"Cannot derive torchtitan model name from HF model_type '{model_type}'. "
	f"Supported types: {list(_HF_MODEL_TYPE_TO_TORCHTITAN_NAME.keys())}."
	)

	# Import the model package and find the configs dict
	model_module = importlib.import_module(f"torchtitan.models.{name}")
	model_configs = None
	for attr in dir(model_module):
	obj = getattr(model_module, attr)
	if isinstance(obj, dict) and attr.endswith("_configs"):
	model_configs = obj
	break

	if model_configs is None:
	raise ValueError(
	f"Could not find model configs dict in torchtitan.models.{name}. "
	f"Expected a dict attribute ending with '_configs'."
	)

	hidden_size = hf_config.hidden_size
	num_layers = hf_config.num_hidden_layers
	vocab_size = hf_config.vocab_size

	for flavor_name, model_cfg in model_configs.items():
	if (
	getattr(model_cfg, "dim", None) == hidden_size
	and getattr(model_cfg, "n_layers", None) == num_layers
	and getattr(model_cfg, "vocab_size", None) == vocab_size
	):
	logger.info(
	f"Auto-derived torchtitan name='{name}', flavor='{flavor_name}' from HF model_type='{model_type}'"
	)
	return name, flavor_name

	raise ValueError(
	f"No matching torchtitan flavor found for model_type='{model_type}' "
	f"(hidden_size={hidden_size}, num_hidden_layers={num_layers}, "
	f"vocab_size={vocab_size}). "
	f"Available flavors for '{name}': {list(model_configs.keys())}."
	)


	def enable_fsdp_gradient_division(model: nn.Module, dp_size: int) -> None:
	"""
	Re-enable FSDP's automatic gradient division.

	TorchTitan calls disable_fsdp_gradient_division() which sets gradient_divide_factor=1.0.
	This re-enables it by setting the factor to the specified dp_size, so gradients are
	averaged across FSDP ranks. This is needed for verl's loss scaling (loss * dp_size)
	to work correctly.

	Args:
	model: The model (or model part) to enable gradient division on.
	dp_size: The data parallel size to use as the gradient divide factor.
	"""

	for module in model.modules():
	if isinstance(module, FSDPModule):
	module.set_gradient_divide_factor(float(dp_size))


	def get_attention_masks(
	input_batch: torch.Tensor,
	positions: torch.Tensor,
	attn_type: str,
	) -> AttentionMasksType:
	match attn_type:
	case "flex":
	return _get_flex_attention_masks(
	input_batch,
	positions,
	)
	case "varlen":
	return _create_varlen_metadata_for_document(
	input_batch,
	positions,
	)
	case _:
	raise TypeError("Only varlen and flex attn masks are supported")


	def _get_document_mask_mod(positions: torch.Tensor) -> _mask_mod_signature:
	# Detect boundaries from position resets
	first_dummy_value = positions[:, :1] - 1
	position_diff = torch.diff(positions, prepend=first_dummy_value, dim=-1)
	sequence_indices = (position_diff != 1).cumsum(-1) # [batch, seq]

	def document_mask(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor) -> torch.Tensor:
	return sequence_indices[b, q_idx] == sequence_indices[b, kv_idx]

	return document_mask


	def _get_flex_attention_masks(
	input_batch: torch.Tensor,
	positions: torch.Tensor,
	) -> AttentionMasksType:
	mask_mods = [get_causal_mask_mod()]
	B = input_batch.shape[0]
	mask_mods.append(_get_document_mask_mod(positions=positions))
	return create_attention_mask(and_masks(*mask_mods), B, None, input_batch.shape[1], input_batch.shape[1])


	def _create_varlen_metadata_for_document(input_batch: torch.Tensor, positions: torch.Tensor) -> VarlenMetadata:
	"""
	Creates cumulative sequence length indices needed for variable length attention

	Args:
	input_batch: Input token IDs with shape [batch, seq].
	positions: Position IDs with shape [batch, seq]. Boundaries detected where
	position diff != 1 (i.e., position resets).

	Returns:
	VarlenMetadata containing cumulative sequence length indices for q, k, and max_seq_len
	"""
	batch_size, seq_len = input_batch.shape
	device = input_batch.device

	# Detect boundaries from position resets (where diff != 1)
	first_dummy_value = positions[:, :1] - 1
	position_diff = torch.diff(positions, prepend=first_dummy_value, dim=-1)
	# boundary_mask[b, i] is True if position i starts a new document
	boundary_mask = position_diff != 1 # [batch, seq]
	boundary_mask[:, 0] = True

	cu_seqlens_list, all_seq_lengths = [], []
	offset = 0

	for b in range(batch_size):
	# Find positions where new documents start
	boundary_positions = boundary_mask[b].nonzero(as_tuple=True)[0].to(torch.int32)
	sample_cu_seqlens = torch.cat(
	[
	boundary_positions,
	torch.tensor([seq_len], dtype=torch.int32, device=device),
	]
	)
	sample_cu_seqlens = torch.unique_consecutive(sample_cu_seqlens)

	seq_lengths = torch.diff(sample_cu_seqlens)
	all_seq_lengths.append(seq_lengths)

	cu_seqlens_adjusted = sample_cu_seqlens[:-1] + offset
	cu_seqlens_list.append(cu_seqlens_adjusted)

	offset += seq_len

	packed_cu_seqlens = torch.cat(cu_seqlens_list + [torch.tensor([offset], dtype=torch.int32, device=device)])

	max_seqlen = 0
	if len(all_seq_lengths) > 0:
	all_seq_lengths = torch.cat(all_seq_lengths)
	# device to host sync but only done once per model forward
	max_seqlen = all_seq_lengths.max().item()

	return VarlenMetadata(
	cu_seq_q=packed_cu_seqlens,
	cu_seq_k=packed_cu_seqlens,
	max_q=max_seqlen,
	max_k=max_seqlen,
	)


	# Regex to parse: model.layers.{L}.mlp.experts.{E}.{weight_suffix}
	_EXPERT_PATTERN = re.compile(r"\.layers\.(\d+)\..\.experts\.(\d+)\.(.)")


	def _parse_expert_name(name: str) -> tuple[int, int, str] \| None:
	"""Parse layer_id, expert_id, weight_suffix from expert param name."""
	match = _EXPERT_PATTERN.search(name)
	if match:
	return int(match.group(1)), int(match.group(2)), match.group(3)
	return None


	def _make_expert_name_template(name: str) -> str:
	"""Convert 'model.layers.0.mlp.experts.3.w1' -> 'model.layers.0.mlp.experts.{}.w1'"""
	return _EXPERT_PATTERN.sub(lambda m: f".layers.{m.group(1)}.mlp.experts.{{}}.{m.group(3)}", name)


	def iter_per_tensor_params_ep(
	params: dict[str, Any],
	device: int,
	ep_group: torch.distributed.ProcessGroup,
	ep_size: int,
	) -> Generator[tuple[str, torch.Tensor], None, None]:
	"""Yield (name, tensor) pairs for weight sync with Expert Parallel.

	Gathers expert weights across EP ranks one (layer, weight_type) group
	at a time to avoid OOM from materializing all experts simultaneously.

	Non-expert params are yielded first (with FSDP full_tensor() if needed),
	then expert params are all-gathered per group and yielded individually.

	Args:
	params: HF-format state dict with per-expert keys. Expert keys must
	follow the pattern ``model.layers.{L}.mlp.experts.{E}.{suffix}``.
	device: device ID to place tensors on.
	ep_group: The EP process group for all-gather.
	ep_size: Number of EP ranks.
	"""
	expert_params: dict[tuple[int, str], dict[int, tuple[str, Any]]] = defaultdict(dict)
	non_expert_params: list[tuple[str, Any]] = []

	for name, param in params.items():
	parsed = _parse_expert_name(name) if "mlp.experts." in name else None
	if parsed is None:
	non_expert_params.append((name, param))
	else:
	layer_id, expert_id, weight_suffix = parsed
	expert_params[(layer_id, weight_suffix)][expert_id] = (name, param)

	params.clear()

	# Yield non-expert params
	for name, param in non_expert_params:
	if isinstance(param, DTensor):
	yield name, param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True)
	else:
	yield name, param
	del non_expert_params

	# Yield expert params with all-gather
	for (layer_id, weight_suffix), experts_dict in sorted(expert_params.items()):
	sorted_expert_ids = sorted(experts_dict.keys())

	# Stack local expert weights
	local_weights = []
	for eid in sorted_expert_ids:
	_, param = experts_dict[eid]
	if isinstance(param, DTensor):
	param = param.to(device, non_blocking=True).full_tensor()
	else:
	param = param.to(device, non_blocking=True)
	local_weights.append(param)

	name_template = _make_expert_name_template(experts_dict[sorted_expert_ids[0]][0])
	local_stacked = torch.stack(local_weights, dim=0)

	# All-gather across EP ranks
	gathered_list = [torch.empty_like(local_stacked) for _ in range(ep_size)]
	torch.distributed.all_gather(gathered_list, local_stacked, group=ep_group)
	all_experts = torch.cat(gathered_list, dim=0)

	for expert_id in range(all_experts.shape[0]):
	yield name_template.format(expert_id), all_experts[expert_id].to(torch.bfloat16).clone()

	del local_weights, local_stacked, gathered_list, all_experts
	torch.cuda.empty_cache()