[Fix] fix te version compatibility (#25)

**Fix TransformerEngine version string compatibility**

`get_te_version()` failed to correctly parse TE version strings in the
format `0.1.0+te2.9.0`, which is common when TE is installed from
non-standard sources. `PkgVersion` treats `+te2.9.0` as a local version
identifier, causing version comparisons like `>= 2.9` to silently
produce wrong results.

**Changes:**

Added a `parse_te_version_str()` helper that uses a regex to extract the
actual version number from the `+te` suffix (e.g. `0.1.0+te2.9.0` →
`2.9.0`) before passing it to `PkgVersion`, ensuring version-gated code
paths behave correctly.

**Scope:** `megatron/core/utils.py` — TE version detection only, no
functional behavior change.

**Raw error log** 
> [default0]:[rank0]: Traceback (most recent call last):
[default5]:[rank5]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/model_provider.py",
line 67, in model_provider
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/train_gpt.py",
line 247, in <module>
[default3]:[rank3]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/transformer_engine/pytorch/cpu_offload.py",
line 700, in get_cpu_offload_context
[default0]:[rank0]:     pretrain(
[default5]:[rank5]: return model_builder(args, pre_process,
post_process, vp_stage, config=config, pg_collection=pg_collection)
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/training/training.py",
line 911, in pretrain
[default5]:[rank5]:
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[default0]:[rank0]: model, optimizer, opt_param_scheduler =
setup_model_and_optimizer(
[default3]:[rank3]:     raise ValueError(
[default0]:[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/training/training.py",
line 1423, in setup_model_and_optimizer
[default3]:[rank3]: ValueError: CPU Offloading is enabled while it is
not mentioned what to offload (weights/activations)
[default0]:[rank0]: model = get_model(model_provider_func, model_type)
[default5]:[rank5]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/gpt_builders.py",
line 88, in gpt_builder
[default0]:[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/training/training.py",
line 1193, in get_model
[default0]:[rank0]:     model = build_model()
[default0]:[rank0]:             ^^^^^^^^^^^^^
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/training/training.py",
line 1185, in build_model
[default5]:[rank5]:     model = GPTModel(
[default5]:[rank5]:             ^^^^^^^^^
[default0]:[rank0]: model = model_provider_func(pre_process=pre_process,
post_process=post_process)
[default5]:[rank5]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/megatron/core/models/gpt/gpt_model.py",
line 204, in __init__
[default5]:[rank5]:     self.decoder = TransformerBlock(
[default5]:[rank5]:                    ^^^^^^^^^^^^^^^^^
[default5]:[rank5]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/megatron/core/transformer/transformer_block.py",
line 321, in __init__
[default5]:[rank5]:     get_cpu_offload_context(
[default0]:[rank0]:
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[default5]:[rank5]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/megatron/core/extensions/transformer_engine.py",
line 1928, in get_cpu_offload_context
[default5]:[rank5]:     context, sync_func = _get_cpu_offload_context(
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/model_provider.py",
line 67, in model_provider
[default5]:[rank5]:                          ^^^^^^^^^^^^^^^^^^^^^^^^^
[default5]:[rank5]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/transformer_engine/pytorch/cpu_offload.py",
line 700, in get_cpu_offload_context
[default0]:[rank0]: return model_builder(args, pre_process,
post_process, vp_stage, config=config, pg_collection=pg_collection)
[default5]:[rank5]:     raise ValueError(
[default5]:[rank5]: ValueError: CPU Offloading is enabled while it is
not mentioned what to offload (weights/activations)
[default0]:[rank0]:
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/flagscale/train/megatron/gpt_builders.py",
line 88, in gpt_builder
[default0]:[rank0]:     model = GPTModel(
[default0]:[rank0]:             ^^^^^^^^^
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/megatron/core/models/gpt/gpt_model.py",
line 204, in __init__
[default0]:[rank0]:     self.decoder = TransformerBlock(
[default0]:[rank0]:                    ^^^^^^^^^^^^^^^^^
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/megatron/core/transformer/transformer_block.py",
line 321, in __init__
[default0]:[rank0]:     get_cpu_offload_context(
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/megatron/core/extensions/transformer_engine.py",
line 1928, in get_cpu_offload_context
[default0]:[rank0]:     context, sync_func = _get_cpu_offload_context(
[default0]:[rank0]:                          ^^^^^^^^^^^^^^^^^^^^^^^^^
[default0]:[rank0]: File
"/root/miniconda3/envs/flagscale-train/lib/python3.12/site-packages/transformer_engine/pytorch/cpu_offload.py",
line 700, in get_cpu_offload_context
[default0]:[rank0]:     raise ValueError(
[default0]:[rank0]: ValueError: CPU Offloading is enabled while it is
not mentioned what to offload (weights/activations)

Author: Darryl233

megatron/core/package_info.py

@@ -2,15 +2,16 @@
 
 
 MAJOR = 0
-MINOR = 16
+MINOR = 1
 PATCH = 0
-PRE_RELEASE = 'rc0'
+PRE_RELEASE = ''
+LOCAL = 'megatron0.15.0rc7'
 
 # Use the following formatting: (major, minor, patch, pre-release)
 VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)
 
 __shortversion__ = '.'.join(map(str, VERSION[:3]))
-__version__ = '.'.join(map(str, VERSION[:3])) + ''.join(VERSION[3:])
+__version__ = '.'.join(map(str, VERSION[:3])) + ''.join(VERSION[3:]) + (f'+{LOCAL}' if LOCAL else '')
 
 __package_name__ = 'megatron_core'
 __contact_names__ = 'NVIDIA'

megatron/core/utils.py

@@ -327,9 +327,18 @@ def get_te_version_str():
         else:
             return version("transformer-engine")
 
+    def parse_te_version_str(ver_str):
+        import re
+
+        # Handle versions like "0.1.0+te2.9.0" — extract the part after "+te"
+        match = re.search(r'\+te(\d+\.\d+.*)', ver_str)
+        if match:
+            return match.group(1)
+        return ver_str
+
     global _te_version
     if _te_version is None and HAVE_TE:
-        _te_version = PkgVersion(get_te_version_str())
+        _te_version = PkgVersion(parse_te_version_str(get_te_version_str()))
     return _te_version