Repository: EleutherAI/gpt-neox

Issue:
MoE loss variable not defined in gpt j residual code path
Running pythia 14M on master:

```
File "/gpt-neox/train.py", line 34, in <module>
    main()
  File "/gpt-neox/train.py", line 30, in main
    pretrain(neox_args=neox_args)
  File "/gpt-neox/megatron/training.py", line 228, in pretrain
    iteration = train(
  File "/gpt-neox/megatron/training.py", line 913, in train
    loss_dict, skipped_iter = train_step(
  File "/gpt-neox/megatron/training.py", line 793, in train_step
    loss = forward_step(
  File "/gpt-neox/megatron/training.py", line 391, in forward_step
    maybe_tuple = model((tokens, position_ids, attention_mask), neox_args=neox_args)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/engine.py", line 1822, in forward
    loss = self.module(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/gpt-neox/megatron/model/utils.py", line 190, in forward
    x = func(forward_input)
  File "/gpt-neox/megatron/model/utils.py", line 181, in exec_func
    inputs = layer(inputs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/gpt-neox/megatron/model/transformer.py", line 1167, in forward
    output, moe_loss = super().forward(hidden_states, attention_mask)
  File "/gpt-neox/megatron/model/transformer.py", line 1155, in forward
    return output, moe_loss
```

The `moe_loss` variable is not defined inside the `forward()` for `ParallelTransformerLayer`, [in the gpt-j-residual==true branch](https://github.com/EleutherAI/gpt-neox/blob/86758c350bff4a06beefe859b4a27a5bc930facf/megatron/model/transformer.py#L1049)

The variable was introduced when #1129 was merged. I am not too familiar with MoE, maybe @yang can comment on this?


Discussion:

Resolving patch:
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
--- a/megatron/model/transformer.py
+++ b/megatron/model/transformer.py
@@ -1046,6 +1046,7 @@ def _get_bias_dropout(self):
     def forward(self, x, attention_mask, layer_past=None):
         layer_past = layer_past if layer_past is not None else self.layer_past
         bias_dropout_fn = self._get_bias_dropout()
+        moe_loss = torch.tensor(0.0, device=x.device, dtype=x.dtype)
         # x: [b, s, h]
         if self.gpt_j_residual:
             # pseudocode:
@@ -1127,9 +1128,6 @@ def forward(self, x, attention_mask, layer_past=None):
 
             # output = x + mlp(ln2(x))
             layernorm_output = self.post_attention_layernorm(attention_output)
-            moe_loss = torch.tensor(
-                0.0, device=layernorm_output.device, dtype=layernorm_output.dtype
-            )
             mlp_bias = torch.tensor(
                 0.0, device=layernorm_output.device, dtype=layernorm_output.dtype
             )


Test patch: