FlagScale Train Upgrade: Synchronization with Megatron-LM Core0.17.0 by lxd-cumt · Pull Request #1193 · flagos-ai/FlagScale

lxd-cumt · 2026-04-21T08:03:40Z

FlagScale Training Upgrade

Sync FlagScale training layer with upstream Megatron-LM core_r0.17.0. Key changes:

Replace training/ and legacy/ with upstream 0.17.0 content, then re-apply FlagScale patches
Unify tokenizer system with registration-based factory dispatch, remove legacy_tokenizer flag
Replace torch.cuda calls with cur_platform abstraction in training/models
Sync arguments.py with upstream FP4 param gather changes, guard hetero-specific args
Fix loss accumulator reset bug (halved lm_loss at iter 2)
Remove FlagScale-forked Muon optimizer code, use upstream 0.17.0 implementation
Remove stale moe_use_legacy_grouped_gemm references
Found that test_train_robobrain_x0.py was leaking MagicMock objects into sys.modules during collection, breaking test_parallel_context.py. Applied the fix by restoring modules immediately after import instead of deferring cleanup.

….15.0rc7 → core_v0.16.1) Resolved 66+ merge conflicts across 9 files in flagscale/train/megatron/: - training.py (22 conflicts): _STARTUP_TIMESTAMPS, TECudaGraphHelper, RL sequence packing, cur_platform stream - arguments.py (24 conflicts): print_rank_0, CUDA graph expansion, hetero guards, FP8/FP4 args - checkpointing.py (9 conflicts): dp_cp_group, distributed quant summary, convert_to_ep - initialize.py (4 conflicts): set_global_writers, cur_platform stream, hetero CPU comm, dualpipev - utils.py (7 conflicts): cu_seqlens, max_seqlen, local_cp_size, cur_platform - global_vars.py: disable_jit_fuser + set_global_writers - dist_signal_handler.py: SIGNAL_MAP dict, string-based signal constructor - biencoder_dataset_utils.py: cur_platform import - __init__.py: clean merge Preserved all FlagScale customizations: cur_platform abstractions, hetero support, dualpipev pipeline parallelism, engram embedding, spiky loss detection, extra validation datasets, FlagScale marker blocks.

…device() in training.py and utils.py Stage 5 multi-platform compliance: converted 12 torch.cuda.current_device() calls and 1 hardcoded device='cuda' to cur_platform equivalents. Remaining torch.cuda calls are CUDA-specific profiler/memory APIs with no platform abstraction.

…form compliance in train_engram.py - Add lazy import for StableLM2SchedulerConfig at usage site to avoid circular import - Replace 5 torch.cuda.current_device() + 5 .cuda() calls in train_engram.py with cur_platform Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Add FlagScale-specific files that were missing from dev-train: - arguments_fs.py, extra_valid.py, fs_theoretical_memory_usage.py - spiky_loss.py, stablelm2_scheduler.py, peft/ - datasets/concated_indexed_dataset.py, datasets/sft_dataset_fs.py - tokenizer/rwkv_tokenization.py, tokenizer/tokenization_utils.py - SYNC_REPORT.md

…dd_argument calls Upstream core_v0.16.1 moved many manual add_argument calls to ArgumentGroupFactory (TransformerConfig, TrainingConfig, LoggerConfig, etc.). The three-way merge kept the old manual calls, causing argparse.ArgumentError: conflicting option string at runtime. Rebuilt arguments.py from upstream and re-applied FlagScale customizations: - imports: arguments_fs, platform abstraction - hetero parallelism support in validate_args - FFN hidden size with multiple_of/hidden_dim_multiplier - cur_platform.get_device_capability() for GroupedGEMM - engram config support - flagcx distributed backend - FlagScale tokenizer types Moved FlagScale-only arg (inference-wandb-logging-step-interval) to arguments_fs.py. Removed 6 args from arguments_fs.py that are now auto-generated by ArgumentGroupFactory(TransformerConfig): cpu-offloading-num-layers, fp8/fp4-quantizer-factory, moe-flex-dispatcher-backend, moe-hybridep-num-sms, moe-router-padding-for-quantization. Also removed duplicate muon-momentum (now in upstream).

…from arguments_fs.py These args are now auto-generated by ArgumentGroupFactory(TransformerConfig) via inheritance from ModelParallelConfig. Manual add_argument calls caused argparse.ArgumentError at runtime.

…ry exclude list Prevents duplicate CLI argument generation for fields now defined in TransformerConfig: use_dualpipev, moe_fb_overlap, te_fl_prefer, qk_layernorm_hidden_dim, recompute_*_per_stage_micro_batch, peft_type, lora_* fields.

…ramGroupOverride get_standard_config_overrides() returns ParamGroupOverride dicts which are incompatible with _get_param_groups() that expects OptimizerConfig dataclass instances. Reverted to old-style config_overrides construction using copy.deepcopy(config) for decoupled_lr support.

get_batch_on_this_tp_rank now returns 8 keys (added cu_seqlens, max_seqlen, local_cp_size) but forward_step only needs 5. Return explicit keys instead of batch.values().

All training YAMLs that had eval_iters but no eval_interval would crash with 'unsupported operand type(s) for //: int and NoneType' at training.py:3838 because args.eval_interval defaults to None. Fixed across: aquila, deepseek_v3, llava_onevision, qwen2_5, qwen2_5_vl, qwen3, qwen3_vl, qwq, robobrain_x0

…7.0 content

Three-way merge of FlagScale patches (from pre-replacement commit e0c9233) onto upstream 0.17.0 base for all training/ and legacy/model/ files. Key merge decisions: - arguments.py: keep FlagScale args (hetero, auto_tune, flagcx, etc.) + upstream new args - checkpointing.py: keep platform plugin + upstream deletion process tracking - initialize.py: keep pre/post validate_args, spiky_loss; drop fused_kernels (removed in 0.17.0); use upstream init_persistent_async_worker(args.rank, 'forkserver') signature - training.py: combine auto_tune guard + upstream skip_train check; update cuda graph API names - utils.py: drop old cu_seqlens broadcast (removed upstream); keep optimizer config + has_nvrx - common_config.py: add flagcx to distributed_backend Literal - tokenizer/: restore FlagScale-patched tokenizer files (removed in 0.17.0, still needed) - legacy/model/: clean merges for all 7 files

…0.17.0 - Remove legacy_tokenizer: true from all training YAML configs - Update tokenizer imports: use megatron.training.tokenizer instead of megatron.core.tokenizers.text.utils - Simplify tokenizer logic: remove legacy_tokenizer branch in train_gpt/train_engram - Delete deprecated tokenizer files: bert_tokenization, multimodal_tokenizer, sft_tokenizer, tokenization_utils - Refactor tokenizer.py to be a thin FlagScale extension layer over upstream build_tokenizer - Update auto_tuner and global_vars imports

…0.17.0 - Remove old get_megatron_optimizer_config from utils.py (used MuonOptimizerConfig, shadowed by the upstream-aligned version in training.py) - Remove MuonOptimizerConfig/AdamOptimizerConfig/SGDOptimizerConfig/ParamKey imports from utils.py (not in upstream) - Remove _add_regularization_args from arguments_fs.py (--muon-matched-adamw-rms, --muon-ns-steps, --no-muon-nesterov duplicated/conflicted with upstream args) - Remove stale 'from megatron.training.utils import get_megatron_optimizer_config' from training.py Muon optimizer now fully uses upstream core 0.17.0 implementation via megatron.core.optimizer.muon with args defined in arguments.py.

…re_v0.17.0 - gpt_builders.py: add experimental_attention_variant handling, fix MTP logic to use get_gpt_decoder_layer_specs, preserve FlagScale hetero support - model_provider.py: update OOM observer to use torch.cuda.memory._dump_snapshot - train_gpt.py: full sync with upstream pretrain_gpt.py (packed seq, MTP, hybrid CP, FIM, startup timing, get_embedding_ranks) - train_aquila_sft.py, train_engram.py, train_qwen2_5_vl.py, train_qwen3_vl.py, train_robobrain_x0.py, train_rwkv.py: replace stale modelopt_args_enabled() with getattr(args, 'modelopt_enabled', False), add missing pretrain import

…actory dispatch - Rewrite tokenizer.py with _TOKENIZER_FACTORY_REGISTRY and register_tokenizer_factory() - Each FlagScale tokenizer type owns its special token definitions internally - build_tokenizer checks FlagScale registry first, falls through to upstream - Add unique_identifiers property to _FlagScaleTokenizerBase for dataset serialization - Update train_gpt.py to import build_tokenizer from FlagScale's shim - Verified: Qwen3-32b training completes 30 iterations successfully

…loss at iter 2 The upstream is_first_iteration flag sets should_reset=False for iteration 1, which prevents the advanced_iters counter from resetting. However, the loss tensor values were unconditionally zeroed after printing. This mismatch caused iteration 2 to report actual_loss/2 because the counter was 2 but only one iteration's loss was accumulated. Now both the loss tensor zeroing and the counter reset are conditional on should_reset, keeping them in sync.

…ipeline_layer_split The dataclass types in ModelParallelConfig (str for enable_hetero, bare list for hetero_pipeline_layer_split) are incompatible with ArgumentGroupFactory auto-generation. Add them to the exclude list and restore the original manual definitions (store_true for enable_hetero, nargs=* type=int for hetero_pipeline_layer_split). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…for hetero mode In validate_args, data_parallel_size and expert_tensor_parallel_size were being unconditionally recalculated from mesh-local TP/PP values, overwriting the correct values set by pre_validate_args in hetero mode. Guard both assignments with 'if not enable_hetero'. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

- engram_builder.py: remove moe_use_legacy_grouped_gemm param and args - train_aquila_sft.py: remove moe_use_legacy_grouped_gemm from layer spec calls - arguments_fs.py: use getattr for moe_use_legacy_grouped_gemm assertion

…le training/models Replace hardcoded torch.cuda.* calls with megatron.plugin.platform abstraction (cur_platform) for multi-platform support (CUDA, NPU, MUSA). Files modified: - flagscale/train/processor/device_processor.py (is_available) - flagscale/train/utils/random_utils.py (is_available, get/set_rng_state) - flagscale/train/megatron/training/training.py (use_mem_pool, empty_cache) - flagscale/train/train_pi.py (current_device) - flagscale/models/megatron/engram/engram.py (Stream, stream, Event, etc.) - flagscale/models/megatron/engram/engram_model.py (Stream, stream, etc.) - flagscale/models/megatron/engram/multi_head_embedding.py (current_device, nvtx) - flagscale/models/megatron/qwen3_vl/language_model.py (current_device) - flagscale/models/megatron/qwen2_5_vl/tensor_parallel.py (current_device) Skipped: CUDA debug internals (memory._snapshot, check_error/cudart), docstring-only references, and commented-out lines.

aoyulong · 2026-04-21T11:23:09Z

Should we delete this file?

GroupedMLP was removed in upstream core_r0.17.0 (replaced by TEGroupedMLP), so moe_use_legacy_grouped_gemm no longer exists as an arg. The getattr fallback always returned False, making this assertion dead code.

…anges - Replace fp4_param with fp4_param_gather and use graceful warning instead of hard error - Add kw_args['fp4_param'] = args.fp4_param_gather in core_transformer_config_from_args - Add --fp4-param-gather CLI argument - Reorganize TE args section (FP4/FP8/precision config) - Add fp4_param to deprecated config list

zhaoyinglia · 2026-04-22T01:54:07Z

+######### FlagScale Begin #########
+# Use FlagScale's tokenizer shim which handles FlagScale-specific tokenizer
+# types (QwenTokenizerFS, etc.) and delegates standard types to upstream.
+from megatron.training.tokenizer import build_tokenizer


move from megatron.training.tokenizer import build_tokenizer to line 65, put “flagscale import” togather

zhaoyinglia · 2026-04-22T02:32:46Z

+        # In RL inference-only mode, train_iters must still be set despite having no optimizer.
+        if args.perform_rl_step:
+            update_train_iters(args)
+    else:
        config, config_overrides = get_megatron_optimizer_config(args)


add if para_ctx is not None: banch

LiJunscs · 2026-04-22T06:15:24Z

All files of engram have no problem. LGTM

…v-train-0.17

legitnull

LGTM

heavyrain-lzy · 2026-04-23T02:45:54Z

-        self.video_token = '<|video_pad|>'
-        self.vision_start_token = '<|vision_start|>'
-        self.vision_end_token = '<|vision_end|>'
+class _Qwen2VLTokenizer(_FlagScaleTokenizerBase):


Only need to replace the class _Qwen2VLTokenizer(MegatronLegacyTokenizer): using class _Qwen2VLTokenizer(_FlagScaleTokenizerBase):. Otherwise, an error will be reported.

…v-train-0.17

Without eval_interval, training.py crashes with TypeError when computing eval_samples: args.train_iters // args.eval_interval where eval_interval is None.

lxd-cumt and others added 22 commits April 17, 2026 16:30

dev-train: replace training/legacy with upstream core_v0.16.1 content

4d7b669

fix(train): remove --enable-hetero and --hetero-pipeline-layer-split …

9256d43

…from arguments_fs.py These args are now auto-generated by ArgumentGroupFactory(TransformerConfig) via inheritance from ModelParallelConfig. Manual add_argument calls caused argparse.ArgumentError at runtime.

fix(train_gpt): unpack only 5 values from get_batch

31fbaa2

get_batch_on_this_tp_rank now returns 8 keys (added cu_seqlens, max_seqlen, local_cp_size) but forward_step only needs 5. Return explicit keys instead of batch.values().

dev-train-0.17: replace training/ and legacy/ with upstream core_v0.1…

cb0bf0d

…7.0 content

lxd-cumt requested review from aoyulong, heavyrain-lzy and zhaoyinglia as code owners April 21, 2026 08:03

lxd-cumt changed the title ~~FlagScale Trainb Upgrade: Synchronization with Megatron-LM Core0.17.0~~ FlagScale Train Upgrade: Synchronization with Megatron-LM Core0.17.0 Apr 21, 2026

aoyulong reviewed Apr 21, 2026

View reviewed changes

lxd-cumt added 2 commits April 22, 2026 10:31

fix: remove stale moe_use_legacy_grouped_gemm check from arguments_fs.py

0cb3f9f

GroupedMLP was removed in upstream core_r0.17.0 (replaced by TEGroupedMLP), so moe_use_legacy_grouped_gemm no longer exists as an arg. The getattr fallback always returned False, making this assertion dead code.

zhaoyinglia reviewed Apr 22, 2026

View reviewed changes

LiJunscs reviewed Apr 22, 2026

View reviewed changes

Comment thread flagscale/train/megatron/train_engram.py

LiJunscs Apr 22, 2026

Copy link
Copy Markdown

Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All files of engram have no problem. LGTM

lxd-cumt force-pushed the dev-train-0.17 branch from 49f837c to 02cb070 Compare April 22, 2026 07:37

lxd-cumt added 7 commits April 22, 2026 15:40

Merge branch 'main' of https://github.com/flagos-ai/FlagScale into de…

6ac5a06

…v-train-0.17

remove aquila

06b7f65

remove detect_all_factory_conflicts.py and SYNC_REPORT.md from FlagScale

6da2c0a

Merge branch 'main' of https://github.com/flagos-ai/FlagScale into de…

de8711d

…v-train-0.17

format

a01aa13

tmp change Megatron-LM-FL repo source for cicd testing

6e3eb68

fix para_ctx

42f45a9

legitnull previously approved these changes Apr 23, 2026

View reviewed changes

heavyrain-lzy reviewed Apr 23, 2026

View reviewed changes

Refactor Qwen2VLTokenizer to use AutoTokenizer

2085492

heavyrain-lzy dismissed legitnull’s stale review via 2085492 April 23, 2026 03:33

lxd-cumt added 12 commits April 23, 2026 16:29

from gitee to github, for cicd test

af6d74b

fix

580cd9c

fix null tokenizer

4fa89c3

code format

c9572e9

fix mock errors

1ba6d33

Merge branch 'main' of https://github.com/flagos-ai/FlagScale into de…

cef9b26

…v-train-0.17

fix: add missing eval_interval to engram CI/CD test config

ecca73b

Without eval_interval, training.py crashes with TypeError when computing eval_samples: args.train_iters // args.eval_interval where eval_interval is None.

fix: add missing eval_interval to deepseek tp2_pp2_ep2 CI/CD test config

22c3cd9

fix

d7cf9ad

tmp cancel update functional test logs

056eefb

update golden values

620a8e2

reset cicd config

5b29b10

aoyulong approved these changes May 6, 2026

View reviewed changes

zhaoyinglia approved these changes May 6, 2026

View reviewed changes

aoyulong merged commit 09e1ca1 into flagos-ai:main May 6, 2026
84 of 105 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

FlagScale Train Upgrade: Synchronization with Megatron-LM Core0.17.0#1193

FlagScale Train Upgrade: Synchronization with Megatron-LM Core0.17.0#1193
aoyulong merged 44 commits into
flagos-ai:mainfrom
lxd-cumt:dev-train-0.17

lxd-cumt commented Apr 21, 2026 •

edited

Loading

Uh oh!

aoyulong Apr 21, 2026

Uh oh!

lxd-cumt Apr 22, 2026

Uh oh!

zhaoyinglia Apr 22, 2026

Uh oh!

zhaoyinglia Apr 22, 2026

Uh oh!

lxd-cumt Apr 22, 2026

Uh oh!

LiJunscs Apr 22, 2026

Uh oh!

legitnull left a comment

Uh oh!

heavyrain-lzy Apr 23, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

lxd-cumt commented Apr 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

aoyulong Apr 21, 2026

Choose a reason for hiding this comment

Uh oh!

lxd-cumt Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

zhaoyinglia Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

zhaoyinglia Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

lxd-cumt Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

LiJunscs Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

legitnull left a comment

Choose a reason for hiding this comment

Uh oh!

heavyrain-lzy Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

lxd-cumt commented Apr 21, 2026 •

edited

Loading

heavyrain-lzy Apr 23, 2026 •

edited

Loading