Error while using aggregate_moe_loss_stats

[prathamk-tw]

When I use aggregate_moe_loss_stats function then i get import error

Traceback (most recent call last):
File "/nemo_run/code/train_moe.py", line 141, in _log_moe_metrics
moe_loss_dict = aggregate_moe_loss_stats(loss_scale=1.0)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 1833, in aggregate_moe_loss_stats
tracker = parallel_state.get_moe_layer_wise_logging_tracker()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'megatron.core.parallel_state' has no attribute 'get_moe_layer_wise_logging_tracker'

This is present in nvcr.io/nvidia/nemo:25.07 & nvcr.io/nvidia/nemo:25.07.02

[prathamk-tw]

Created a PR for this #14694

[zhenyih]

ai generated, please verify

The error occurs because the function get_moe_layer_wise_logging_tracker is located in the wrong module. In NeMo containers 25.07 and 25.07.02, the aggregate_moe_loss_stats function is trying to access this function from megatron.core.parallel_state, but it's actually defined in megatron.core.transformer.moe.moe_utils.

To fix this issue, modify your code to import the function from the correct location:

# Add this import at the top of your file
from megatron.core.transformer.moe.moe_utils import get_moe_layer_wise_logging_tracker

# Then create a wrapper function that uses the correct import
def aggregate_moe_loss_stats_fixed(loss_scale=1.0):
    # Use the imported function instead of parallel_state.get_moe_layer_wise_logging_tracker()
    tracker = get_moe_layer_wise_logging_tracker()
    # Rest of your function remains the same
    # ...

Use this wrapper function instead of the original aggregate_moe_loss_stats to avoid the AttributeError.

[prathamk-tw]

Yes this is what i have done the PR above. Please review it and merge. Thanks