question of PP x aux_loss for MoE

[rakkit]

In short, does PP allow multiple-args input and multiple-args output?

——

Hey, we’ve been stuck for a while on how to properly integrate aux loss for MoE training with PP and compile(full_graph).

For context, both DeepSeek V3 and GLM 4.5 mention that

“We also applied an auxiliary sequence-level balance loss with a 0.0001 weight to avoid extreme imbalance within any single sequence.”

(We could open a PR for the sequence-level balance loss if you’re interested.)

To make this work, we need to compute the extra loss at each block, either by:

• Caching the per-layer aux_loss loss (which breaks compile, but not PP), or

• Passing both activations and aux_loss to the next PP stage (which doesn’t affect compile).

The second option basically requires the PP API to support multiple-args input and output. We tried earlier this year to explicitly pass arguments when building PP stages, but it didn’t work. I’m wondering if there have been any updates since then, or if we might have missed something.

Do you have any other suggestions or better solutions? @tianyu-l @H-Huang
CC: @janEbert @garrett361

[garrett361]

We tried earlier this year to explicitly pass arguments when building PP stages, but it didn’t work

What happened, exactly?

[rakkit]

@garrett361 by default, PP's Stage infers the output shape automatically, and we tried to manually tell the input/output shapes for multiple inputs/outputs when building the PP Stage. Then got a bunch of errors 🥲 and it's just not working

[garrett361]

Ah, yeah, I recall having to explicitly pass in the shapes for some previous PP work where the auto-inference code was erroring out for me (for reasons I forget now). I eventually got it working, but one issue I remember encountering was just getting the dtypes wrong for a while and that being hard to debug. Sounds like your errors are more complicated, though.

Would be helpful if you could post the specific errors you're hitting, though, if you can.

[tianyu-l]

@rakkit Yes I was thinking about adding an aux loss. It'd be great if you could contribute a PR. cc @shuhuayu

Could you share the specific aux loss definition you are using? Sharing some pseudocode on how you use PP with it would also be helpful.

cc @xmfan @wconstab

[janEbert]

Similar to gradient accumulation, we wrap the loss function to support PP:

def moe_aux_loss(
    pred: MoEInputsDict,
    labels: torch.Tensor,
    loss_fn: LossFunction,
) -> torch.Tensor:
    """Sequence-wise auxiliary loss-enhanced loss function for MoE Transformer
    model training.
    """
    loss = loss_fn(pred, labels)
    if isinstance(pred, dict):
        loss += pred["aux_loss"]
    return loss

In train.py Trainer.__init__:

        self.loss_fn = self.train_spec.build_loss_fn(job_config)

        # `MoETransformer` does not exist in upstream, would need to adapt the if-query.
        if issubclass(self.train_spec.model_cls, MoETransformer):
            pre_moe_aux_loss_fn = self.loss_fn
            self.loss_fn = functools.partial(
                moe_aux_loss,
                loss_fn=pre_moe_aux_loss_fn,
            )

[shuhuayu]

@rakkit Thanks for bringing this up. Two questions here:

1. Do you detach with the backbone so the aux loss is only used to update routers in current MoE layers?
2. Can we do backward per pipeline stage so we do not need to pass aux losses to the final task loss?

[rakkit]

Thanks @shuhuayu, i did not get your first question. And for 2, yes, we can, and it can avoid pp's multi-in, multi-out, but it does not work well with the compiler

[shuhuayu]

Thanks @shuhuayu, i did not get your first question. And for 2, yes, we can, and it can avoid pp's multi-in, multi-out, but it does not work well with the compiler

Thanks for the clarification. My question 1 is that when doing backward from aux loss, do you backward through all weights before router module or detaching at router module, in the second case gradients regarding to aux loss are only computed on routers.

[rakkit]

Thanks @shuhuayu, I see. That's a good question. I might be wrong, but I have not seen people detach the graph and prevent each layer’s auxiliary loss from propagating to previous layers. In our internal tests, we allow all auxiliary losses to backpropagate to the inputs.

(I think detaching would also break the compiler and the full graph, since we would need to add a gradient hook to handle it.)

[shuhuayu]

Thanks for the clarifications @rakkit ! It would be great if you could contribute a PR and we can look into it.

[man2machine]

@rakkit Yes, its possible to have multiple inputs and outputs for PP, and I have tested this and it works even in PyTorch versions from a while back (I tried this back in PyTorch 2.6 and it worked). But there are some caveats. You can take a look at these issues for more info:
pytorch/pytorch#152827
pytorch/pytorch#159711

So far, I found that if you output or input a tuple of tensors, with the batch dimension being the 0th axis on each tensor, PP works with no problems. You can look at the PipelineStage source code to see how this is supported.

If you want to implement a load balancing loss for MoE as described in various papers in an easy way, then you can just output multiple tensors in a PP stage, and accumulate the losses until the last layer. This is what I do and it works just fine. It might not be the most efficient though for all losses, since if the loss does not propagate gradients all the way to the input, it might introduce an additional communication overhead.

If you want a loss function to update weights in a particular pipeline stage only based on inputs for only that stage, then I would suggest not using pipeline parallelism for that loss, and instead just having a separate gradient update outside of the pipeline parallel stuff. PP already utilizes a separate optimizer for each stage anyways, so you can just introduce another optimizer for that loss. PP in pytorch implements a particular set of algorithms which assume that the loss is computed from the outputs on the last stage.

If you can show the error messages or give a minimal reproducible example then I can see if this is a new problem, or something that is already supported and if you are just missing something.

[rakkit]

Thanks a lot for the useful discussion, i opened the PR here.

it surprisingly works well with 1F1B scheduler at the latest Torch (2.10.0.dev20251118+cu129), but was broken for other PP schedulers.

[man2machine]

By the way @rakkit I meant that PyTorch PP supports multiple inputs and outputs and was not referring to TorchTitan. As for whether TorchTitan will error if you do that because it has some additional logic outside of the PyTorch library, I don't know. So a PR to support that on TorchTitan as you are doing might still be needed. But the actual PyTorch library definitely does support multiple inputs/outputs.