Unified multiple platform support for flagscale training (flagos-ai#1146)

## Integrates with the Megatron-LM-FL multi-platform plugin system

Reference: 
- Platform cuda support, for nvidia chips or other cuda-compatible
chips: flagos-ai/Megatron-LM-FL#11
- Platform musa support, for mthreads chips:
flagos-ai/Megatron-LM-FL#14

---------

Co-authored-by: zihugithub <fbye@baai.ac.cn>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 3 people

.github/workflows/functional_tests_hetero_train.yml

@@ -166,6 +166,7 @@ jobs:
             ${INSTALL_DIR:+--install-dir "$INSTALL_DIR"} \
             --no-system --no-dev --no-base --no-task \
             --src-deps megatron-lm \
+            --force-build \
             --retry-count 3
 
           # Install FlagScale CLI

.github/workflows/functional_tests_train.yml

@@ -155,8 +155,24 @@ jobs:
           command -v flagscale || { echo "❌ FlagScale CLI not found in PATH"; exit 1; }
           echo "✅ FlagScale CLI installed successfully: $(flagscale --version 2>/dev/null || echo 'version unknown')"
 
-          # For train task: all dependencies are pre-installed in the env
-          # No additional installation needed
+          # Install Megatron-LM-FL from source (force-build to replace pre-installed megatron-core)
+          # Derive install-dir from env_path (e.g., /root/miniconda3 -> /root)
+          INSTALL_DIR=""
+          if [ "$PKG_MGR" = "conda" ] && [ -n "$ENV_PATH" ]; then
+            INSTALL_DIR=$(dirname "$ENV_PATH")
+          fi
+
+          ./tools/install/install.sh \
+            --platform ${{ inputs.platform }} \
+            --task train \
+            --pkg-mgr "$PKG_MGR" \
+            ${ENV_NAME:+--env-name "$ENV_NAME"} \
+            ${INSTALL_DIR:+--install-dir "$INSTALL_DIR"} \
+            --no-system --no-dev --no-base --no-task \
+            --src-deps megatron-lm \
+            --force-build \
+            --retry-count 3
+
           echo "Environment ready for train tests"
         timeout-minutes: 30
 

.github/workflows/unit_tests_common.yml

@@ -154,6 +154,7 @@ jobs:
           fi
 
           # Only install Megatron-LM source dep (pip deps are pre-installed in Docker image)
+          # Use --force-build to ensure Megatron-LM-FL replaces any pre-installed megatron-core
           ./tools/install/install.sh \
             --platform ${{ inputs.platform }} \
             --task train \
@@ -163,6 +164,7 @@ jobs:
             --no-system --no-dev --no-base --no-task \
             --src-deps megatron-lm \
             --pip-deps typer \
+            --force-build \
             --retry-count 3
 
           # TODO: temp solution to install newly added deps, remove once the new image is built

flagscale/runner/launcher/launcher_ssh.py

@@ -345,7 +345,10 @@ def run(
         if enable_monitoring is None:
             enable_monitoring = self.config.experiment.runner.get("enable_monitoring", False)
         num_visible_devices = None
-        visible_devices = self.user_envs.get("CUDA_VISIBLE_DEVICES", None)
+        # visible_devices = self.user_envs.get("CUDA_VISIBLE_DEVICES", None)
+        visible_devices = next(
+            (v for k, v in self.user_envs.items() if k.endswith("_VISIBLE_DEVICES")), None
+        )
         if visible_devices is not None and isinstance(visible_devices, str):
             visible_devices = visible_devices.split(",")
             num_visible_devices = len(visible_devices)
@@ -983,7 +986,10 @@ def _run_gpu_health_check(self):
 
                 # Get CUDA_VISIBLE_DEVICES if set
                 cur_envs = add_decive_extra_config(self.user_envs, resource_info["type"])
-                visible_devices = cur_envs.get("CUDA_VISIBLE_DEVICES", None)
+                # visible_devices = cur_envs.get("CUDA_VISIBLE_DEVICES", None)
+                visible_devices = next(
+                    (v for k, v in cur_envs.items() if k.endswith("_VISIBLE_DEVICES")), None
+                )
                 num_visible_devices = None
                 if visible_devices is not None and isinstance(visible_devices, str):
                     visible_devices = visible_devices.split(",")
@@ -1050,7 +1056,10 @@ def run_health_check_thread(node_config):
             node_rank = 0
             host = "localhost"
 
-            visible_devices = self.user_envs.get("CUDA_VISIBLE_DEVICES", None)
+            # visible_devices = self.user_envs.get("CUDA_VISIBLE_DEVICES", None)
+            visible_devices = next(
+                (v for k, v in self.user_envs.items() if k.endswith("_VISIBLE_DEVICES")), None
+            )
             num_visible_devices = None
             if visible_devices is not None and isinstance(visible_devices, str):
                 visible_devices = visible_devices.split(",")

flagscale/train/megatron/legacy/data/biencoder_dataset_utils.py

@@ -13,6 +13,9 @@
 from megatron.training import get_args, get_tokenizer, print_rank_0
 from megatron.training.datasets.data_samplers import MegatronPretrainingSampler
 
+from megatron.plugin.platform import get_platform
+cur_platform = get_platform()
+
 
 def make_attention_mask(source_block, target_block):
     """
@@ -192,7 +195,7 @@ def get_block_samples_mapping(block_dataset, title_dataset, data_prefix, num_epo
     # This should be a barrier but nccl barrier assumes
     # device_index=rank which is not the case for model
     # parallel case
-    counts = torch.tensor([1], dtype=torch.long, device='cuda')
+    counts = torch.tensor([1], dtype=torch.long, device=cur_platform.device_name())
     torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
     assert counts[0].item() == mpu.get_data_parallel_group().size()
 

flagscale/train/megatron/legacy/data/dataset_utils.py

@@ -34,6 +34,8 @@
 from megatron.core import mpu
 from megatron.core.datasets.indexed_dataset import IndexedDataset
 
+from megatron.plugin.platform import get_platform
+cur_platform = get_platform()
 
 DSET_TYPE_BERT = 'standard_bert'
 DSET_TYPE_ICT = 'ict'
@@ -707,7 +709,7 @@ def get_samples_mapping(indexed_dataset,
     # This should be a barrier but nccl barrier assumes
     # device_index=rank which is not the case for model
     # parallel case
-    counts = torch.tensor([1], dtype=torch.long, device='cuda')
+    counts = torch.tensor([1], dtype=torch.long, device=cur_platform.device_name())
     torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
     torch.distributed.all_reduce(counts, group=mpu.get_pipeline_model_parallel_group())
     assert counts[0].item() == (

flagscale/train/megatron/legacy/data/realm_dataset_utils.py

@@ -10,6 +10,9 @@
 from megatron.legacy.data.dataset_utils import create_masked_lm_predictions, pad_and_convert_to_numpy
 from megatron.training import get_args, get_tokenizer, print_rank_0
 
+from megatron.plugin.platform import get_platform
+cur_platform = get_platform()
+
 
 def get_one_epoch_dataloader(dataset, micro_batch_size=None):
     """Specifically one epoch to be used in an indexing job."""
@@ -179,7 +182,7 @@ def get_block_samples_mapping(block_dataset, title_dataset, data_prefix, num_epo
     # This should be a barrier but nccl barrier assumes
     # device_index=rank which is not the case for model
     # parallel case
-    counts = torch.tensor([1], dtype=torch.long, device='cuda')
+    counts = torch.tensor([1], dtype=torch.long, device=cur_platform.device_name())
     torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
     assert counts[0].item() == mpu.get_data_parallel_group().size()
 

flagscale/train/megatron/legacy/fused_kernels/tests/test_fused_kernels.py

@@ -10,6 +10,9 @@
 from megatron.legacy.model.utils import attention_mask_func
 from megatron.legacy.fused_kernels import load
 
+from megatron.plugin.platform import get_platform
+cur_platform = get_platform()
+
 def test_load_fused_kernels():
     try:
         import fused_layer_norm_cuda
@@ -23,7 +26,7 @@ def test_load_fused_kernels():
         raise e
 
 def test_fused_softmax():
-    bert = BertModel.from_pretrained("bert-base-cased").cuda().half()
+    bert = BertModel.from_pretrained("bert-base-cased").to(cur_platform.device()).half()
     tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
     test_text = (
         "Hello. How are you? I am fine thank you and you? yes Good. "
@@ -36,16 +39,16 @@ def test_fused_softmax():
     )
 
     embedding_output = bert.embeddings(
-        input_ids=tokens["input_ids"].cuda(),
+        input_ids=tokens["input_ids"].to(cur_platform.device()),
         position_ids=None,
-        token_type_ids=tokens["token_type_ids"].cuda(),
+        token_type_ids=tokens["token_type_ids"].to(cur_platform.device()),
         inputs_embeds=None,
         past_key_values_length=0,
     )
 
     # (bsz, 1, 1, seq_len)
     mask = bert.get_extended_attention_mask(
-        attention_mask=tokens["attention_mask"].cuda(),
+        attention_mask=tokens["attention_mask"].to(cur_platform.device()),
         input_shape=tokens["input_ids"].shape,
         device=bert.device,
     )
@@ -69,7 +72,7 @@ def test_fused_softmax():
             attn_mask_type=AttnMaskType.padding,
             scaled_masked_softmax_fusion=True,
         )
-        .cuda()
+        .to(cur_platform.device())
         .half()
     )
 
@@ -88,7 +91,7 @@ def test_fused_softmax():
             attn_mask_type=AttnMaskType.padding,
             scaled_masked_softmax_fusion=False,
         )
-        .cuda()
+        .to(cur_platform.device())
         .half()
     )
 
@@ -121,7 +124,7 @@ def test_fused_softmax():
 
 
 def test_fused_upper_triangle_mask_softmax():
-    gpt = GPT2Model.from_pretrained("gpt2").cuda().half()
+    gpt = GPT2Model.from_pretrained("gpt2").to(cur_platform.device()).half()
     tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
     test_text = (
         "Hello. How are you? I am fine thank you and you? yes Good. "
@@ -133,14 +136,14 @@ def test_fused_upper_triangle_mask_softmax():
         return_tensors="pt",
     )
 
-    attention_mask = tokens["attention_mask"].cuda()
+    attention_mask = tokens["attention_mask"].to(cur_platform.device())
     attention_mask = attention_mask.view(attention_mask.size(0), -1)
     attention_mask = attention_mask[:, None, None, :]
     attention_mask = (1.0 - attention_mask) * -10000.0
     attention_mask = attention_mask.repeat(1, 1, attention_mask.size()[-1], 1)
     attn = gpt.h[0]
 
-    hidden_states = gpt.wte(tokens["input_ids"].cuda())
+    hidden_states = gpt.wte(tokens["input_ids"].to(cur_platform.device()))
     q, k, v = attn.attn.c_attn(hidden_states).split(768, dim=-1)
     q = attn.attn._split_heads(q, attn.attn.num_heads, attn.attn.head_dim)
     k = attn.attn._split_heads(k, attn.attn.num_heads, attn.attn.head_dim)
@@ -169,7 +172,7 @@ def test_fused_upper_triangle_mask_softmax():
             attn_mask_type=AttnMaskType.causal,
             scaled_masked_softmax_fusion=True,
         )
-        .cuda()
+        .to(cur_platform.device())
         .half()
     )
 
@@ -188,7 +191,7 @@ def test_fused_upper_triangle_mask_softmax():
             attn_mask_type=AttnMaskType.causal,
             scaled_masked_softmax_fusion=False,
         )
-        .cuda()
+        .to(cur_platform.device())
         .half()
     )
 
@@ -221,7 +224,7 @@ def test_fused_upper_triangle_mask_softmax():
 
 
 def test_layer_norm():
-    bert = BertModel.from_pretrained("bert-base-cased").cuda().half()
+    bert = BertModel.from_pretrained("bert-base-cased").to(cur_platform.device()).half()
     tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
     test_text = (
         "Hello. How are you? I am fine thank you and you? yes Good. "
@@ -236,22 +239,22 @@ def test_layer_norm():
     # [bsz, seq_len, d_model]
     embedding_output = (
         bert.embeddings(
-            input_ids=tokens["input_ids"].cuda(),
+            input_ids=tokens["input_ids"].to(cur_platform.device()),
             position_ids=None,
-            token_type_ids=tokens["token_type_ids"].cuda(),
+            token_type_ids=tokens["token_type_ids"].to(cur_platform.device()),
             inputs_embeds=None,
             past_key_values_length=0,
         )
-        .cuda()
+        .to(cur_platform.device())
         .half()
     )
 
     fused_layernorm_layer = (
-        MixedFusedLayerNorm(normalized_shape=embedding_output.size(-1)).cuda().half()
+        MixedFusedLayerNorm(normalized_shape=embedding_output.size(-1)).to(cur_platform.device()).half()
     )
 
     torch_layernorm_layer = (
-        LayerNorm(normalized_shape=embedding_output.size(-1)).cuda().half()
+        LayerNorm(normalized_shape=embedding_output.size(-1)).to(cur_platform.device()).half()
     )
 
     fused_output = fused_layernorm_layer(embedding_output)
@@ -299,8 +302,8 @@ def test_masked_softmax_forward():
     scale_t = torch.tensor([1.0])
     for qlen in [128, 256, 1024, 2048, 4096]:
         for klen in [128, 256, 1024, 2048]:
-            inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0')
-            masks = torch.randint(0, 2, (batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0')
+            inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device=cur_platform.device_name(0))
+            masks = torch.randint(0, 2, (batch, 1, qlen, klen), dtype=torch.bool, device=cur_platform.device_name(0))
             softmax_results = scaled_masked_softmax_cuda.forward(inputs, masks, scale_t[0].item())
             softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t[0].item())
             error = (softmax_results_torch - softmax_results).abs().max()
@@ -314,9 +317,9 @@ def test_masked_softmax_backward():
     scale_t = torch.tensor([1.0])
     for qlen in [128, 256, 1024, 2048, 4096]:
         for klen in [128, 256, 1024, 2048]:
-            inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0')
-            backward = torch.rand_like(inputs, dtype=torch.float16, device='cuda:0')
-            masks = torch.randint(0, 2, (batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0')
+            inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device=cur_platform.device_name(0))
+            backward = torch.rand_like(inputs, dtype=torch.float16, device=cur_platform.device_name(0))
+            masks = torch.randint(0, 2, (batch, 1, qlen, klen), dtype=torch.bool, device=cur_platform.device_name(0))
             softmax_results = scaled_masked_softmax_cuda.forward(inputs, masks, scale_t[0].item())
             back_grad = scaled_masked_softmax_cuda.backward(backward, softmax_results, scale_t[0].item())
 
@@ -335,8 +338,8 @@ def test_allmasked_softmax_forward():
     scale_t = torch.tensor([1.0])
     for qlen in [128, 256, 1024, 2048, 4096]:
         for klen in [128, 256, 1024, 2048]:
-            inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0')
-            masks = torch.ones((batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0')
+            inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device=cur_platform.device_name(0))
+            masks = torch.ones((batch, 1, qlen, klen), dtype=torch.bool, device=cur_platform.device_name(0))
             softmax_results = scaled_masked_softmax_cuda.forward(inputs, masks, scale_t[0].item())
             softmax_results_torch = torch.zeros_like(inputs)
             error = (softmax_results_torch - softmax_results).abs().max()
@@ -351,9 +354,9 @@ def test_allmasked_softmax_backward():
     scale_t = torch.tensor([1.0])
     for qlen in [128, 256, 1024, 2048, 4096]:
         for klen in [128, 256, 1024, 2048]:
-            inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0')
-            backward = torch.rand_like(inputs, dtype=torch.float16, device='cuda:0')
-            masks = torch.ones((batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0')
+            inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device=cur_platform.device_name(0))
+            backward = torch.rand_like(inputs, dtype=torch.float16, device=cur_platform.device_name(0))
+            masks = torch.ones((batch, 1, qlen, klen), dtype=torch.bool, device=cur_platform.device_name(0))
             softmax_results = scaled_masked_softmax_cuda.forward(inputs, masks, scale_t[0].item())
             back_grad = scaled_masked_softmax_cuda.backward(backward, softmax_results, scale_t[0].item())
             inputs.requires_grad = True

flagscale/train/megatron/legacy/model/bert_model.py

@@ -16,6 +16,8 @@
 from megatron.legacy.model.utils import scaled_init_method_normal
 from .module import MegatronModule
 
+from megatron.plugin.platform import get_platform
+cur_platform = get_platform()
 
 def bert_extended_attention_mask(attention_mask):
     # We create a 3D attention mask from a 2D tensor mask.
@@ -200,7 +202,7 @@ def forward(self, bert_model_input, attention_mask,
                 output = torch.zeros(
                     size=(embeddings.shape[0], embeddings.shape[2]),
                     dtype=torch.float32,
-                    device=torch.cuda.current_device())
+                    device=cur_platform.current_device())
                 for i, (embedding, mask) in enumerate(zip(embeddings, masks)):
                     output[i, :] = torch.mean(embedding[1: mask - 1], dim=0)
 

flagscale/train/megatron/legacy/model/module.py

@@ -10,9 +10,12 @@
 from megatron.core import mpu, tensor_parallel
 
 
-_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
-_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
-_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor)
+from megatron.plugin.platform import get_platform
+cur_platform = get_platform()
+
+_FLOAT_TYPES = (torch.FloatTensor, cur_platform.FloatTensor)
+_HALF_TYPES = (torch.HalfTensor, cur_platform.HalfTensor)
+_BF16_TYPES = (torch.BFloat16Tensor, cur_platform.BFloat16Tensor)
 
 
 def param_is_not_shared(param):
@@ -117,7 +120,7 @@ def initialize_word_embeddings(self):
         # values.
         if mpu.is_rank_in_embedding_group(ignore_virtual=False):
             self.shared_embedding_or_output_weight().data = (
-                self.shared_embedding_or_output_weight().data.cuda()
+                self.shared_embedding_or_output_weight().data.to(cur_platform.device())
             )
             torch.distributed.all_reduce(
                 self.shared_embedding_or_output_weight().data, group=mpu.get_embedding_group()