Akoumparouli/nemo ux moe loss logging
State: closed | Pull request: yes
Labels: name: Run CICD

# What does this PR do ?
Without this PR

```
Training epoch 0, iteration 0/9 | lr: 0.0006 | consumed_samples: 32 | global_batch_size: 32 | global_step: 0 | reduced_train_loss: 11.02
Training epoch 0, iteration 1/9 | lr: 0.0006 | consumed_samples: 64 | global_batch_size: 32 | global_step: 1 | reduced_train_loss: 9.761
Training epoch 0, iteration 2/9 | lr: 0.0006 | consumed_samples: 96 | global_batch_size: 32 | global_step: 2 | reduced_train_loss: 9.21
Training epoch 0, iteration 3/9 | lr: 0.0006 | consumed_samples: 128 | global_batch_size: 32 | global_step: 3 | reduced_train_loss: 8.595
Training epoch 0, iteration 4/9 | lr: 0.0006 | consumed_samples: 160 | global_batch_size: 32 | global_step: 4 | reduced_train_loss: 8.675
Training epoch 0, iteration 5/9 | lr: 0.0006 | consumed_samples: 192 | global_batch_size: 32 | global_step: 5 | reduced_train_loss: 7.918
Training epoch 0, iteration 6/9 | lr: 0.0006 | consumed_samples: 224 | global_batch_size: 32 | global_step: 6 | reduced_train_loss: 7.319
Training epoch 0, iteration 7/9 | lr: 0.0006 | consumed_samples: 256 | global_batch_size: 32 | global_step: 7 | reduced_train_loss: 7.889
Training epoch 0, iteration 8/9 | lr: 0.0006 | consumed_samples: 288 | global_batch_size: 32 | global_step: 8 | reduced_train_loss: 7.113
Training epoch 0, iteration 9/9 | lr: 0.0006 | consumed_samples: 320 | global_batch_size: 32 | global_step: 9 | reduced_train_loss: 6.909
```


With this PR:

```
Training epoch 0, iteration 0/9 | lr: 0.0006 | consumed_samples: 32 | global_batch_size: 32 | global_step: 0 | reduced_train_loss: 11.02 | load_balancing_loss: 9.812
Training epoch 0, iteration 1/9 | lr: 0.0006 | consumed_samples: 64 | global_batch_size: 32 | global_step: 1 | reduced_train_loss: 9.763 | load_balancing_loss: 9.919
Training epoch 0, iteration 2/9 | lr: 0.0006 | consumed_samples: 96 | global_batch_size: 32 | global_step: 2 | reduced_train_loss: 9.232 | load_balancing_loss: 13.04
Training epoch 0, iteration 3/9 | lr: 0.0006 | consumed_samples: 128 | global_batch_size: 32 | global_step: 3 | reduced_train_loss: 8.579 | load_balancing_loss: 16.77
Training epoch 0, iteration 4/9 | lr: 0.0006 | consumed_samples: 160 | global_batch_size: 32 | global_step: 4 | reduced_train_loss: 8.661 | load_balancing_loss: 17.5
Training epoch 0, iteration 5/9 | lr: 0.0006 | consumed_samples: 192 | global_batch_size: 32 | global_step: 5 | reduced_train_loss: 7.909 | load_balancing_loss: 13.16
Training epoch 0, iteration 6/9 | lr: 0.0006 | consumed_samples: 224 | global_batch_size: 32 | global_step: 6 | reduced_train_loss: 7.358 | load_balancing_loss: 13.96
Training epoch 0, iteration 7/9 | lr: 0.0006 | consumed_samples: 256 | global_batch_size: 32 | global_step: 7 | reduced_train_loss: 7.994 | load_balancing_loss: 14.74
Training epoch 0, iteration 8/9 | lr: 0.0006 | consumed_samples: 288 | global_batch_size: 32 | global_step: 8 | reduced_train_loss: 6.796 | load_balancing_loss: 15.49
Training epoch 0, iteration 9/9 | lr: 0.0006 | consumed_samples: 320 | global_batch_size: 32 | global_step: 9 | reduced_train_loss: 6.789 | load_balancing_loss: 16.16
```

```
config = llm.MixtralConfig8x7B(
        num_layers=2,
        hidden_size=768,
        ffn_hidden_size=3072,
        num_attention_heads=8,
        seq_length=seq_length,
        init_method_std=0.023,
        hidden_dropout=0.1,
        attention_dropout=0.1,
        layernorm_epsilon=1e-5,
        make_vocab_size_divisible_by=128,
        masked_softmax_fusion=args.masked_softmax_fusion,
        # Moe,
        num_moe_experts=8,
        moe_aux_loss_coeff=0.01,
        moe_router_topk=2,
        moe_router_pre_softmax=True,
        moe_token_dispatcher_type="alltoall",
        moe_router_load_balancing_type='aux_loss',
    )
```
**Collection**: [Note which collection this PR will affect]

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage
* You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

This PR was closed because it has been inactive for 7 days since being marked as stale.

[🤖]: Hi @akoumpa 👋,

We wanted to let you know that a [CICD pipeline](https://github.com/NVIDIA/NeMo/actions/runs/11334331014) for this PR just finished successfully

So it might be time to merge this PR or get some approvals

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g