Update TE-FL arguments (flagos-ai#972)

### PR Category
Train

### PR Types
Bug Fixes

### PR Description
Redesign the parameters in TE-FL as follows:

- **use_transformer_engine_fl = True**
- Enables TE-FL and, by default, also enables flag_gems.enable. All
Torch Ops in the framework-side Megatron-LM are replaced with Gems Ops,
while the TransformerEngine side enters the new-implementation branch
and likewise replaces every Torch Op with a Gems Op.
  
- **use_transformer_engine_fl = True && disable_flag_gems_replace_torch
= True**
- Enables TE-FL but does NOT enable flag_gems.enable. All Torch Ops in
the framework-side Megatron-LM keep their original implementations, and
the TransformerEngine side still enters the new-implementation branch
yet retains every original Torch Op implementation.
- This mode is useful for debugging the training process to determine
whether an issue lies in the operator implementation or in the TE-FL
plugin-style implementation.
  
- **use_flag_gems_replace_torch = True**
- Does NOT enable TE-FL, but does enable flag_gems.enable. All Torch Ops
in the framework-side Megatron-LM are replaced with Gems Ops, while the
TransformerEngine side continues to use the original FusedOp
implementation and does NOT enter the new branch.

Author: lxd-cumt

flagscale/backends/Megatron-LM/megatron/core/tensor_parallel/layers.py.patch

@@ -1,5 +1,5 @@
 diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py
-index 5c6f34b70..5c943759a 100644
+index 5c6f34b70..c4f77ea61 100644
 --- a/megatron/core/tensor_parallel/layers.py
 +++ b/megatron/core/tensor_parallel/layers.py
 @@ -448,6 +448,7 @@ class LinearWithGradAccumulationAndAsyncCommunication(torch.autograd.Function):
@@ -56,13 +56,13 @@ index 5c6f34b70..5c943759a 100644
              # Need to return None's as gradient has to flow for all the input arguments
              # provided during forward
 -            return (sub_grad_input, grad_weight, grad_bias, None, None, None, None, None, None)
-+            return (sub_grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None, None)
++            return (sub_grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None)
 
          if ctx.allreduce_dgrad:
              handle.wait()
 
 -        return grad_input, grad_weight, grad_bias, None, None, None, None, None, None
-+        return grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None, None
++        return grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None
 
 
  def linear_with_grad_accumulation_and_async_allreduce(

flagscale/backends/Megatron-LM/megatron/training/arguments.py.patch

@@ -1,5 +1,5 @@
 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
-index 1120c7529..761a796db 100644
+index 1120c7529..3facfbf30 100644
 --- a/megatron/training/arguments.py
 +++ b/megatron/training/arguments.py
 @@ -67,6 +67,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser):
@@ -202,7 +202,7 @@ index 1120c7529..761a796db 100644
          else:
              args.ffn_hidden_size = 4 * args.hidden_size
 
-@@ -1175,6 +1198,148 @@ def validate_args(args, defaults={}):
+@@ -1175,6 +1198,147 @@ def validate_args(args, defaults={}):
              args.recompute_granularity != 'full'
          ), 'recompute_granularity must not be full when CUDA Graphs are enabled.'
 
@@ -342,16 +342,15 @@ index 1120c7529..761a796db 100644
 +        assert args.ckpt_format == 'torch', "PEFT is only tested with torch format checkpoint"
 +
 +    if args.use_transformer_engine_fl:
-+        assert not args.add_bias_linear, "Do not support linear with bias in TransformerEngine-TL Now"
-+        assert args.context_parallel_size == 1, "DotProductAttention in TransformerEngine-FL, do not support context parallel now"
 +        assert not args.tp_comm_overlap, "Do not support gemm/sp comm in TransformerEngine-FL now"
 +        args.distributed_backend = 'flagcx'
-+        args.use_flag_gems = True
++        if not args.disable_flag_gems_replace_torch:
++            args.use_flag_gems_replace_torch = True
 +
      # Print arguments.
      _print_args("arguments", args)
 
-@@ -1585,6 +1750,8 @@ def _add_network_size_args(parser):
+@@ -1585,6 +1749,8 @@ def _add_network_size_args(parser):
                         help='Which normalization technique to use.')
      group.add_argument('--norm-epsilon', type=float, default=1e-5,
                         help='Epsilon for layer norm and RMS norm.')
@@ -360,7 +359,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--apply-layernorm-1p', action='store_true',
                         help='Adjust LayerNorm weights such that they are centered '
                         'around zero. This improves numerical stability.')
-@@ -1608,6 +1775,10 @@ def _add_network_size_args(parser):
+@@ -1608,6 +1774,10 @@ def _add_network_size_args(parser):
      group.add_argument('--glu-linear-offset', type=float, default=0.0,
                         help='Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). '
                              'Only used when gated_linear_unit is True')
@@ -371,7 +370,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--onnx-safe', type=bool, required=False,
                         help='Use workarounds for known problems with '
                         'Torch ONNX exporter')
-@@ -1820,6 +1991,14 @@ def _add_logging_args(parser):
+@@ -1820,6 +1990,14 @@ def _add_logging_args(parser):
                         help='The wandb experiment name.')
      group.add_argument('--wandb-save-dir', type=str, default='',
                         help='Path to save the wandb results locally.')
@@ -386,7 +385,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--logging-level', type=int, default=None,
                         help='Set default logging level')
      return parser
-@@ -2001,6 +2180,25 @@ def _add_training_args(parser):
+@@ -2001,6 +2179,25 @@ def _add_training_args(parser):
                         '"shared_experts": recompute the shared experts in the MoE layer.'
                         '"moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, '
                         '"core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing.')
@@ -412,7 +411,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--no-clone-scatter-output-in-embedding', action='store_false',
                         help='If not set, clone the output of the scatter in embedding layer to GC original tensor.',
                         dest='clone_scatter_output_in_embedding')
-@@ -2087,6 +2285,10 @@ def _add_training_args(parser):
+@@ -2087,6 +2284,10 @@ def _add_training_args(parser):
                         help='Total number of samples to train over all '
                         'training runs. Note that either train-iters or '
                         'train-samples should be provided.')
@@ -423,7 +422,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--log-interval', type=int, default=100,
                         help='Report loss and timing interval.')
      group.add_argument('--exit-interval', type=int, default=None,
-@@ -2210,6 +2412,10 @@ def _add_training_args(parser):
+@@ -2210,6 +2411,10 @@ def _add_training_args(parser):
                         help='The communicator group names to use high priority streams.')
      group.add_argument('--use-te-activation-func', action='store_true',
                         help='Use activation function kernel from Transformer Engine in MLP module.')
@@ -434,7 +433,7 @@ index 1120c7529..761a796db 100644
 
      return parser
 
-@@ -2268,11 +2474,26 @@ def _add_learning_rate_args(parser):
+@@ -2268,11 +2473,26 @@ def _add_learning_rate_args(parser):
                         'and initial warmup, the learning rate at each '
                         'iteration would be different.')
      group.add_argument('--lr-decay-style', type=str, default='linear',
@@ -462,7 +461,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--lr-decay-iters', type=int, default=None,
                         help='number of iterations to decay learning rate over,'
                         ' If None defaults to `--train-iters`')
-@@ -2331,6 +2552,8 @@ def _add_checkpointing_args(parser):
+@@ -2331,6 +2551,8 @@ def _add_checkpointing_args(parser):
      group.add_argument('--save-retain-interval', type=int, default=None,
                         help='Number of iterations between retained checkpoints (other'
                         'checkpoints _except the last checkpoint_ are automatically deleted).')
@@ -471,7 +470,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--no-save-optim', action='store_true', default=None,
                         help='Do not save current optimizer.')
      group.add_argument('--no-save-rng', action='store_true', default=None,
-@@ -2380,6 +2603,8 @@ def _add_checkpointing_args(parser):
+@@ -2380,6 +2602,8 @@ def _add_checkpointing_args(parser):
      group.add_argument('--no-use-tokenizer-model-from-checkpoint-args', action='store_false',
                         dest='use_tokenizer_model_from_checkpoint_args',
                         help='If set, do not use tokenizer model path from checkpoint')
@@ -480,7 +479,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--exit-on-missing-checkpoint', action='store_true',
                         help="If '--load' is set, but checkpoint is not found "
                         "(e.g., path typo), then exit instead of random "
-@@ -2541,7 +2766,7 @@ def _add_distributed_args(parser):
+@@ -2541,7 +2765,7 @@ def _add_distributed_args(parser):
                         default=False, help='if set, overlap pipeline parallel communication in warmup and flush',
                         dest='overlap_p2p_comm_warmup_flush')
      group.add_argument('--distributed-backend', default='nccl',
@@ -489,7 +488,7 @@ index 1120c7529..761a796db 100644
                         help='Which backend to use for distributed training.')
      group.add_argument('--distributed-timeout-minutes', type=int, default=10,
                         help='Timeout minutes for torch.distributed.')
-@@ -2592,6 +2817,11 @@ def _add_distributed_args(parser):
+@@ -2592,6 +2816,11 @@ def _add_distributed_args(parser):
                         'complete it instead. Also turns on '
                         '--use-cpu-initialization flag. This is for '
                         'external DDP manager.' )
@@ -501,7 +500,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--account-for-embedding-in-pipeline-split', action='store_true',
                         default=False, help='If set, *input* embedding layer will be treated as a standard transformer'
                         'layer in the context of partition and placement for pipeline parallelism.')
-@@ -2636,6 +2866,10 @@ def _add_distributed_args(parser):
+@@ -2636,6 +2865,10 @@ def _add_distributed_args(parser):
                         help='If set, keep the fp8 transpose cache when using Megatron FSDP.')
      group.add_argument('--enable-full-sharding-in-hsdp', action='store_true',
                         help='If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.')
@@ -512,7 +511,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--num-distributed-optimizer-instances', type=int, default=1,
                         help='Number of Distributed Optimizer copies across Data Parallel domain.')
      group.add_argument('--use-torch-fsdp2', action='store_true',
-@@ -2690,6 +2924,9 @@ def _add_validation_args(parser):
+@@ -2690,6 +2923,9 @@ def _add_validation_args(parser):
      group.add_argument('--eval-interval', type=int, default=1000,
                         help='Interval between running evaluation on '
                         'validation set.')
@@ -522,7 +521,7 @@ index 1120c7529..761a796db 100644
      group.add_argument("--test-mode", action="store_true", help='Run all real-time test alongside the experiment.')
      group.add_argument('--skip-train', action='store_true',
                         default=False, help='If set, bypass the training loop, '
-@@ -2708,6 +2945,8 @@ def _add_tokenizer_args(parser):
+@@ -2708,6 +2944,8 @@ def _add_tokenizer_args(parser):
                         'automatically calculated from vocab-size.')
      group.add_argument('--vocab-file', type=str, default=None,
                         help='Path to the vocab file.')
@@ -531,7 +530,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--merge-file', type=str, default=None,
                         help='Path to the BPE merge file.')
      group.add_argument('--vocab-extra-ids', type=int, default=0,
-@@ -2726,8 +2965,17 @@ def _add_tokenizer_args(parser):
+@@ -2726,8 +2964,17 @@ def _add_tokenizer_args(parser):
                                  'MultimodalTokenizer',
                                  'NullTokenizer',
                                  'NullMultimodalTokenizer',
@@ -550,7 +549,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--tokenizer-model', type=str, default=None,
                         help='Sentencepiece tokenizer model.')
      group.add_argument('--tokenizer-metadata', type=str, default=None,
-@@ -2768,6 +3016,11 @@ def _add_data_args(parser):
+@@ -2768,6 +3015,11 @@ def _add_data_args(parser):
      group.add_argument('--valid-data-path', nargs='*', default=None,
                         help='The weight and prefix list for an independent validation dataset. '
                         'Follows the same pattern rules as --data-path.')
@@ -562,7 +561,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--test-data-path', nargs='*', default=None,
                         help='The weight and prefix list for an independent test dataset. '
                         'Follows the same pattern rules as --data-path.')
-@@ -2816,11 +3069,18 @@ def _add_data_args(parser):
+@@ -2816,11 +3068,18 @@ def _add_data_args(parser):
                         'end-of-document token.')
      group.add_argument('--eod-mask-loss', action='store_true',
                         help='Mask loss for the end of document tokens.')
@@ -581,7 +580,7 @@ index 1120c7529..761a796db 100644
      group.add_argument('--object-storage-cache-path', type=str, default=None,
                         help='Path to cache index files when using s3 or msc dataloader')
      group.add_argument('--mid-level-dataset-surplus', type=float, default=0.005,
-@@ -2897,6 +3157,19 @@ def _add_biencoder_args(parser):
+@@ -2897,6 +3156,19 @@ def _add_biencoder_args(parser):
      return parser
 
 
@@ -601,7 +600,7 @@ index 1120c7529..761a796db 100644
  def _add_vision_args(parser):
      group = parser.add_argument_group(title="vision")
 
-@@ -2967,6 +3240,8 @@ def _add_vision_args(parser):
+@@ -2967,6 +3239,8 @@ def _add_vision_args(parser):
                         help='Whether to layer normalize the q and k attention embeddings.')
      group.add_argument('--qk-l2-norm', action='store_true',
                         help='Use llama 4 qk l2 norm')
@@ -610,7 +609,7 @@ index 1120c7529..761a796db 100644
 
      return parser
 
-@@ -3275,3 +3550,92 @@ def _add_sft_args(parser):
+@@ -3275,3 +3549,94 @@ def _add_sft_args(parser):
      group.add_argument('--sft-tokenizer-prompt-format', type=str, default="nemotron-h-aligned", 
                         help='SFT prompt format.')
      return parser
@@ -690,7 +689,9 @@ index 1120c7529..761a796db 100644
 +    group = parser.add_argument_group(title="flag engine")
 +    group.add_argument('--use-transformer-engine-fl', action='store_true',
 +                       help='Enable transformer engine fl for training')
-+    group.add_argument('--use-flag-gems', action='store_true',
++    group.add_argument('--disable-flag-gems-replace-torch', action='store_true',
++                       help='Disable flag gems for training when using te-fl')
++    group.add_argument('--use-flag-gems-replace-torch', action='store_true',
 +                       help='Enable flag gems for training')
 +    group.add_argument('--flag-gems-log-path', type=str, default=None,
 +                       help='Path of flag gems logging')

flagscale/train/train.py

@@ -2560,7 +2560,7 @@ def get_e2e_base_metrics():
         cuda_graph_helper.create_cudagraphs()
 
     # enable flag_gems for transformer_engine_fl
-    if args.use_flag_gems:
+    if args.use_flag_gems_replace_torch:
         if not HAVE_GEMS:
             raise ValueError(f"Can not import flag gems")
         else: