refit latest update based on ahmads refactor

.gitmodules

@@ -1,6 +1,6 @@
 [submodule "3rdparty/Megatron-LM"]
 	path = 3rdparty/Megatron-LM-workspace/Megatron-LM
-	url = https://github.com/yaoyu-33/Megatron-LM.git
+	url = https://github.com/NVIDIA/Megatron-LM.git
 	branch = main
 	shallow = true
 [submodule "3rdparty/Megatron-Bridge"]

3rdparty/Megatron-LM-workspace/setup.py

@@ -43,7 +43,7 @@
 # VCS dependencies use full "pkg @ git+URL@rev" format matching pyproject.toml [tool.uv.sources]
 CACHED_DEPENDENCIES = [
     # Default dependencies from pyproject.toml
-    "torch",
+    "torch>=2.6.0",
     "numpy",
     "packaging>=24.2",
     # Dev dependencies from pyproject.toml
@@ -58,7 +58,7 @@
     "opentelemetry-api~=1.33.1",
     "mamba-ssm~=2.2",
     "causal-conv1d~=1.5",
-    "flash-linear-attention~=0.3.2",
+    "flash-linear-attention~=0.4.0",
     "nv-grouped-gemm~=1.1",
     "megatron-energon[av_decode]~=6.0",
     "av",
@@ -69,6 +69,9 @@
     "emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.1.0",
     "datasets",
     "fastapi~=0.50",
+    "flask[async]",
+    "hypercorn",
+    "openai",
 ]
 
 

examples/configs/grpo_math_1B.yaml

@@ -146,6 +146,11 @@ policy:
     moe_enable_deepep: false
     moe_token_dispatcher_type: "allgather"
     moe_shared_expert_overlap: false
+    moe_pad_experts_for_cuda_graph_inference: false
+    cuda_graph_impl: "local"
+    cuda_graph_scope: null
+    use_te_rng_tracker: true
+    inference_rng_tracker: true
 
     optimizer:
       optimizer: "adam"
@@ -252,13 +257,16 @@ policy:
     stop_strings: null
     mcore_generation_config:
       buffer_size_gb: 20  # Total GPU memory (in GB) allocated for KV cache buffers
-      buffer_guaranteed_fraction: 0.1  # Fraction of buffer reserved for guaranteed active requests
       num_cuda_graphs: 16  # Number of CUDA graphs to pre-compile for different batch sizes
       block_size_tokens: 256  # Size of each KV cache block in tokens (affects memory granularity)
       use_cuda_graphs_for_non_decode_steps: true  # Enable CUDA graphs for prefill/context processing
-      enable_chunked_prefill: true  # Split long prefills into chunks for better memory management
-      unified_memory_level: 0  # Unified memory usage level (0=disabled, higher values enable more aggressive paging)
+      unified_memory_level: 0  # Unified memory usage level (0=disabled, 1+=enables unified memory )
       max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
+      enable_chunked_prefill: false
+      kv_cache_management_mode: "persist" # Can be "persist", "offload", or "recompute"
+      static_kv_memory_pointers: false # Relevant only for offload and recompute modes
+      materialize_only_last_token_logits: false
+
     vllm_cfg:
       async_engine: false
       precision: ${policy.precision}

examples/configs/grpo_math_1B_megatron.yaml

@@ -100,6 +100,11 @@ policy:
     moe_shared_expert_overlap: false
     #gives ~20% training perf speedup with sequence packing 
     apply_rope_fusion: True
+    moe_pad_experts_for_cuda_graph_inference: false
+    cuda_graph_impl: "local"
+    cuda_graph_scope: null
+    use_te_rng_tracker: true
+    inference_rng_tracker: true
 
     optimizer:
       optimizer: "adam"
@@ -151,9 +156,12 @@ policy:
       num_cuda_graphs: 16  # Number of CUDA graphs to pre-compile for different batch sizes
       block_size_tokens: 256  # Size of each KV cache block in tokens (affects memory granularity)
       use_cuda_graphs_for_non_decode_steps: true  # Enable CUDA graphs for prefill/context processing
-      enable_chunked_prefill: false  # Split long prefills into chunks for better memory management
-      unified_memory_level: 0  # Unified memory usage level (0=disabled, higher values enable more aggressive paging)
+      unified_memory_level: 1  # Unified memory usage level (0=disabled, 1+=enables unified memory )
       max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
+      enable_chunked_prefill: false
+      kv_cache_management_mode: "persist" # Can be "persist", "offload", or "recompute"
+      static_kv_memory_pointers: false # Relevant only for offload and recompute modes
+      materialize_only_last_token_logits: false
 
     vllm_cfg:
       tensor_parallel_size: 1