Skip to content

[Refactor] Remove unused world_size plumbing #1964

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
HAOCHENYE wants to merge 1 commit into
base: reduce-sum-03-switch-sum
Choose a base branch
from
Draft

[Refactor] Remove unused world_size plumbing #1964

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 0 additions & 3 deletions xtuner/v1/loss/aux_loss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,6 @@ def accumulate(
z_ctx: list[ZLossContext] | ZLossContext | None = None,
num_tokens_local: int = 0,
num_tokens_global: torch.Tensor | None = None,
world_size: int = 1,
) -> torch.Tensor:
"""Accumulate routing statistics for one layer and inject z-loss into
the main graph.
Expand All @@ -113,7 +112,6 @@ def accumulate(
num_tokens_global (torch.Tensor | None): All-reduced non-padding token count across
ranks (int64 scalar). Pass ``None`` when ``z_loss_global_average`` is off or no
process group is initialized.
world_size (int): World size that produced ``num_tokens_global``.

Returns:
torch.Tensor: ``hidden_states`` augmented with the per-layer z-loss autograd hook.
Expand All @@ -139,7 +137,6 @@ def accumulate(
router_logits=selected_router_logits,
num_tokens_local=num_tokens_local,
num_tokens_global=num_tokens_global,
world_size=world_size,
)
hidden_states = AuxLossScaler.apply(hidden_states, z_loss_l)

Expand Down
32 changes: 5 additions & 27 deletions xtuner/v1/loss/moe_loss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,32 +5,13 @@
from cyclopts import Parameter
from pydantic import BaseModel, ConfigDict
from torch import distributed as dist
from torch.distributed._functional_collectives import all_reduce

from xtuner.v1.utils.device import get_device


DEVICE = get_device()


class _AllReduce(torch.autograd.Function):
@staticmethod
def forward(ctx, op, group, tensor):
ctx.group = group
ctx.op = op
tensor = tensor.clone(memory_format=torch.contiguous_format)
tensor = all_reduce(tensor, op, group=group)
return tensor

@staticmethod
def backward(ctx, grad_output):
return (None, None) + (_AllReduce.apply(ctx.op, ctx.group, grad_output),)


def all_reduce_autograd(tensor, op, group):
return _AllReduce.apply(op, group, tensor)


class BalancingLossConfig(BaseModel):
"""Balancing loss configuration for MoE models.

Expand Down Expand Up @@ -139,11 +120,11 @@ def finalize(
non_pad_token (int): Number of non-padding tokens on this rank.

Returns:
tuple[torch.Tensor, torch.Tensor]: ``(balancing_loss, local_balancing_loss)``. The first
carries the autograd graph used for backward; in the global-average branch it is reduced
across ranks via ``all_reduce_autograd``. The second is a detached, per-rank local
component for logging whose cross-rank SUM reproduces the first (used by the display
pipeline so that logged curves stay global once backward switches to the local component).
tuple[torch.Tensor, torch.Tensor]: ``(balancing_loss, local_balancing_loss)``. Both are
this rank's per-rank local component (computed from its own ``local_gating_sum`` with the
global detached statistics); the first carries the autograd graph for backward, aggregated
across ranks by the SUM gradient reduction, while the second is detached for the display
pipeline whose cross-rank SUM restores the global balancing loss.
"""
routing_weights_sum_list = self.routing_weights_sum_list
self.routing_weights_sum_list = []
Expand Down Expand Up @@ -260,7 +241,6 @@ def accumulate(
router_logits: torch.Tensor,
num_tokens_local: int,
num_tokens_global: torch.Tensor | None,
world_size: int,
) -> torch.Tensor:
"""Compute z-loss for one layer and return it as a scalar with autograd
attached.
Expand All @@ -278,8 +258,6 @@ def accumulate(
num_tokens_global (torch.Tensor | None): All-reduced non-padding token count across
ranks, as an int64 scalar tensor. ``None`` when ``z_loss_global_average`` is off
or the process group is not initialized.
world_size (int): Number of ranks contributing to ``num_tokens_global``. Ignored when
``num_tokens_global`` is ``None``.

Returns:
torch.Tensor: Per-layer z-loss as a 0-d tensor with autograd graph back to
Expand Down
26 changes: 10 additions & 16 deletions xtuner/v1/model/moe/moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -248,24 +248,23 @@ def _z_loss_dist_token_count(
z_ctx: list[ZLossContext] | ZLossContext | None,
num_tokens_local: int,
device: torch.device | str | int,
) -> tuple[torch.Tensor | None, int]:
) -> torch.Tensor | None:
"""Compute the cross-rank non-padding token count needed by the z-loss
inline path.

Returns ``(num_tokens_global, world_size)``. ``num_tokens_global`` is ``None`` (i.e. skip
global averaging) when there is no z-loss context, when the configured z-loss is not
global-average, or when no process group is initialized.
Returns the global non-padding token count, or ``None`` (i.e. skip global averaging) when
there is no z-loss context, when the configured z-loss is not global-average, or when no
process group is initialized.
"""
if z_ctx is None:
return None, 1
return None
first = z_ctx[0] if isinstance(z_ctx, list) else z_ctx
if not first.loss_cfg.z_loss_global_average or not dist.is_initialized():
return None, 1
return None
n = torch.tensor(num_tokens_local, device=device, dtype=torch.int64)
group = dist.group.WORLD
assert group is not None
n_global = all_reduce(n, "sum", group)
return n_global, dist.get_world_size()
return all_reduce(n, "sum", group)

def _extract_aux_loss_ctx(
self,
Expand Down Expand Up @@ -487,7 +486,7 @@ def _micro_batch_forward(
[{} for _ in range(len(seq_ctx_list))] if keep_router else []
)
balancing_ctx, z_ctx = self._extract_aux_loss_ctx(loss_ctx_list)
num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, cat_mask.device)
num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, cat_mask.device)

# Process through layers
cat_seq_ctx: SequenceContext | None = None
Expand Down Expand Up @@ -569,7 +568,6 @@ def _micro_batch_forward(
z_ctx=z_ctx,
num_tokens_local=non_pad_token,
num_tokens_global=num_tokens_global,
world_size=z_world_size,
)

assert hidden_states_list, "XTuner Internal Error, found empty hidden states for domino EP"
Expand Down Expand Up @@ -754,7 +752,7 @@ def _forward(
# Hoisted out of the per-layer accumulate path: mask is constant across layers.
nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1]
non_pad_token = nonpad_indices.numel()
num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device)
num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device)

for idx, decoder_layer in self.layers.items():
if int(idx) < self.config.first_k_dense_replace:
Expand Down Expand Up @@ -796,7 +794,6 @@ def _forward(
z_ctx=z_ctx,
num_tokens_local=non_pad_token,
num_tokens_global=num_tokens_global,
world_size=z_world_size,
)

if self.config.return_hidden_states:
Expand Down Expand Up @@ -827,9 +824,7 @@ def _forward(
# MTP uses its own mask; main mask's non-pad indices do not apply.
mtp_nonpad_indices = torch.nonzero(mtp_seq_ctx.mask, as_tuple=True)[1]
mtp_non_pad_token = mtp_nonpad_indices.numel()
mtp_num_tokens_global, mtp_z_world_size = self._z_loss_dist_token_count(
z_ctx, mtp_non_pad_token, mtp_seq_ctx.mask.device
)
mtp_num_tokens_global = self._z_loss_dist_token_count(z_ctx, mtp_non_pad_token, mtp_seq_ctx.mask.device)

# Forward through MTP block
mtp_outputs = self.mtp_block(
Expand Down Expand Up @@ -861,7 +856,6 @@ def _forward(
z_ctx=z_ctx,
num_tokens_local=mtp_non_pad_token,
num_tokens_global=mtp_num_tokens_global,
world_size=mtp_z_world_size,
)
mtp_loss, (_, mtp_extra) = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx))
mtp_losses += mtp_loss
Expand Down
3 changes: 1 addition & 2 deletions xtuner/v1/model/moe/qwen3vl_text.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ def _forward(
# Hoisted out of the per-layer accumulate path: mask is constant across layers.
nonpad_indices = torch.nonzero(seq_ctx.mask, as_tuple=True)[1]
non_pad_token = nonpad_indices.numel()
num_tokens_global, z_world_size = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device)
num_tokens_global = self._z_loss_dist_token_count(z_ctx, non_pad_token, seq_ctx.mask.device)

# =====================================================
deepstack_visual_embeds = seq_ctx.deepstack_visual_embeds
Expand Down Expand Up @@ -196,7 +196,6 @@ def _forward(
z_ctx=z_ctx,
num_tokens_local=non_pad_token,
num_tokens_global=num_tokens_global,
world_size=z_world_size,
)

if deepstack_visual_embeds is not None and ((idx := int(idx)) in range(len(deepstack_visual_embeds))):
Expand Down
Loading