[BUG] The logic for calculating the last stage when average loss across microbatches.

[LitLeo]

Describe the bug

After a step is completed, operation ”Average loss across microbatches“ will be performed in the last stage.
For the common case, although setting ignore_virtual=True or False does not affect the logic for calculating the last stage, ignore_virtual=False should be considered more accurate.

Megatron-LM/megatron/training/training.py

Line 847 in 0dd78dd

if mpu.is_pipeline_last_stage(ignore_virtual=True):

if mpu.is_pipeline_last_stage(ignore_virtual=True):
        # Average loss across microbatches.

Expected behavior

if mpu.is_pipeline_last_stage(ignore_virtual=False):
        # Average loss across microbatches.