[tml] Inkling model support

docs/models/thinkingmachines/index.md

@@ -17,7 +17,7 @@ Inkling needs 16 nodes of 4× GB300 and the `radixark/miles:inkling` image:
 
 ```bash
 cd /root/miles
-python scripts/run_inkling_975b.py train \
+python scripts/run_inkling.py train \
    --model-name Inkling --train-mode full --task dapo_math \
    --num-nodes 16 --num-gpus-per-node 4
 ```

docs/models/thinkingmachines/inkling-small.md

@@ -22,18 +22,20 @@ cd /root/miles
 
 # Full-parameter GRPO. 276 B fits with the CPU-offloaded optimizer -
 # no NVMe streaming needed (unlike the 975 B recipe).
-python scripts/run_inkling_975b.py train \
+python scripts/run_inkling.py train \
    --model-name Inkling-Small --train-mode full --task dapo_math \
    --num-nodes 4 --num-gpus-per-node 8 \
+   --lr 5e-5 --rollout-batch-size 64 --global-batch-size 128 \
    --sglang-context-length 4096 --rollout-max-response-len 2048 \
    --extra-args "--offload-train-target cpu --sglang-mem-fraction-static 0.65 \
       --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer"
 
 # LoRA GRPO (rank 32, all-linear), same cluster; adapter-only weight
 # sync swaps in ~4 s per rollout.
-python scripts/run_inkling_975b.py train \
+python scripts/run_inkling.py train \
    --model-name Inkling-Small --train-mode lora --task dapo_math \
    --num-nodes 4 --num-gpus-per-node 8 \
+   --lr 2e-4 --rollout-batch-size 64 --global-batch-size 128 \
    --sglang-context-length 4096 --rollout-max-response-len 2048
 ```
 
@@ -45,4 +47,4 @@ The model definition lives in `scripts/models/inkling-small.sh` (`MODEL_ARGS_NUM
 |---|---|---|---|---|---|---|---|
 | H200 | 32 | 4 | on | 8 | 4 | 1 | `--decoder-last-pipeline-num-layers 7` (42 = 7×5 + 7) |
 
-Batch shape is configurable from the launcher (`--rollout-batch-size`, `--global-batch-size`; defaults 32/64). The validated Small runs used 64/128 with `--lr 6e-6`, producing a steadily rising dapo-math reward curve; LoRA uses the launcher's `lr 5e-6` default.
+Batch shape is configurable from the launcher (`--rollout-batch-size`, `--global-batch-size`; defaults 32/64). The validated Small runs used 64/128: full with `--lr 5e-5`, LoRA with `--lr 2e-4` — both produce a steadily rising dapo-math reward curve. At the launcher's conservative LoRA default (5e-6) the zero-initialised B factors take hundreds of rollouts to accumulate a visible delta-W, which reads as "not learning".

docs/models/thinkingmachines/inkling.md

@@ -32,12 +32,12 @@ docker pull radixark/miles:inkling
 
 # Full-parameter GRPO on 16 nodes x 4 GB300, inside the container
 cd /root/miles
-python scripts/run_inkling_975b.py train \
+python scripts/run_inkling.py train \
    --model-name Inkling --train-mode full --task dapo_math \
    --num-nodes 16 --num-gpus-per-node 4
 
 # LoRA GRPO (rank 32, all-linear), same cluster
-python scripts/run_inkling_975b.py train \
+python scripts/run_inkling.py train \
    --model-name Inkling --train-mode lora --task dapo_math \
    --num-nodes 16 --num-gpus-per-node 4
 ```
@@ -72,12 +72,12 @@ Pass `--hf-checkpoint <path>` to the launcher when the weights are already on a
 
 ### 4.2 HF → Megatron `torch_dist` conversion
 
-Inkling ships in BF16, so conversion is a single distributed `torch_dist` shard (no precision cast). The model definition comes from `scripts/models/inkling-975b.sh`:
+Inkling ships in BF16, so conversion is a single distributed `torch_dist` shard (no precision cast). The model definition comes from `scripts/models/inkling.sh`:
 
 ```bash
 cd /root/miles
-source scripts/models/inkling-975b.sh
-PYTHONPATH=/root/Megatron-LM torchrun \
+source scripts/models/inkling.sh
+CONVERT_KEEP_PP1=1 PYTHONPATH=/root/Megatron-LM torchrun \
    --nproc-per-node 4 --nnodes 4 \
    --master-addr ${MASTER_ADDR} --master-port 12345 \
    --node-rank ${NODE_RANK} \
@@ -90,6 +90,8 @@ PYTHONPATH=/root/Megatron-LM torchrun \
    --save /root/models/Inkling_torch_dist/
 ```
 
+`CONVERT_KEEP_PP1=1` keeps the conversion at PP1: without it the converter auto-bumps PP toward the rank count, which is incompatible with `--tensor-model-parallel-size 4` at 16 ranks.
+
 The saved `torch_dist` checkpoint is parallelism-agnostic: training can load it under any validated TP / PP / EP layout. Point the launcher at the result with `--torch-dist`.
 
 ### 4.3 Multi-node fan-out
@@ -122,7 +124,7 @@ GRPO with truncated importance sampling. The launcher defaults: global batch siz
 
 ### 5.3 Training attention backends
 
-`--inkling-attn-backend` selects the training-side attention implementation:
+The `MILES_INKLING_ATTN_BACKEND` environment variable selects the training-side attention implementation:
 
 | Backend | Role | 8K packed, rel-extent 1024, GB300 |
 |---|---|---|

miles/backends/megatron_utils/actor.py

@@ -52,7 +52,7 @@
 from .ft.in_memory_checkpoint import InMemoryCheckpointManager
 from .ft.indep_dp import reconfigure_indep_dp_group
 from .initialize import init, is_first_replica_megatron_main_rank
-from .lora_utils import is_lora_enabled
+from .lora_utils import is_lora_enabled, lora_rollout_enabled
 from .model import TrainStepOutcome, forward_only, initialize_model_and_optimizer, save, train
 from .parallel import verify_megatron_parallel_state
 from .replay_utils import register_replay_list_moe
@@ -70,6 +70,12 @@
 
 
 def _setup_disk_offload_reclaim(disk_dir: str) -> None:
+    """Wipe this rank's train disk-offload dir on startup and re-arm the atexit wipe.
+
+    torch_memory_saver unlinks each backup file as its allocation is freed on a
+    graceful teardown, but a SIGKILL'd run leaves stale files behind. The dir is
+    per-rank (see actor_factory), so clearing it wholesale touches nobody else.
+    """
     if not disk_dir:
         return
     shutil.rmtree(disk_dir, ignore_errors=True)
@@ -249,7 +255,7 @@ def init(
             weights_getter=lambda: self.weights_backuper.get("actor"),
             model_name=type(self.hf_config).__name__.lower() if self.args.model_name is None else self.args.model_name,
             quantization_config=getattr(self.hf_config, "quantization_config", None),
-            is_lora=is_lora_enabled(args),
+            is_lora=lora_rollout_enabled(args),
         )
 
         # Adapters currently loaded into Megatron slots on this rank.
@@ -289,7 +295,7 @@ def sleep(self) -> None:
 
         destroy_process_groups()
 
-        tag = "default" if is_lora_enabled(self.args) else None
+        tag = "default" if lora_rollout_enabled(self.args) else None
         torch_memory_saver.pause(tag=tag)
 
         self._asleep = True
@@ -308,7 +314,7 @@ def wake_up(self) -> None:
             return
         print_memory("before wake_up model")
 
-        tag = "default" if is_lora_enabled(self.args) else None
+        tag = "default" if lora_rollout_enabled(self.args) else None
         torch_memory_saver.resume(tag=tag)
 
         clear_memory()

miles/backends/megatron_utils/lora_utils.py

@@ -11,7 +11,7 @@
 import torch.distributed as dist
 
 from miles.backends.training_utils.parallel import get_parallel_state
-from miles.utils.lora import is_lora_enabled
+from miles.utils.lora import is_lora_enabled, lora_rollout_enabled  # noqa: F401  (re-exported)
 
 logger = logging.getLogger(__name__)
 
@@ -115,6 +115,60 @@ def lora_base_cpu_backup_enabled(args: Namespace) -> bool:
     return is_lora_enabled(args) and getattr(args, "colocate", False) and getattr(args, "lora_base_cpu_backup", False)
 
 
+def sglang_lora_target_all_sentinel(args) -> bool:
+    """Hand SGLang the ``"all"`` shorthand so it auto-detects module names (required for Inkling)."""
+    from miles.utils.chat_template_utils.inkling import is_inkling_checkpoint
+
+    return is_inkling_checkpoint(getattr(args, "hf_checkpoint", None) or "")
+
+
+_marked_lora_grad_params_cache: dict[int, list] = {}
+
+
+def reduce_marked_lora_grads(model: Sequence[torch.nn.Module]) -> None:
+    """Sum partial grads of replicated LoRA params over their tagged group ("tp"|"ep"), before the DP reduce-scatter."""
+    from megatron.core import parallel_state as ps
+
+    key = id(model[0]) if model else 0
+    marked = _marked_lora_grad_params_cache.get(key)
+    if marked is None:
+        marked = []
+        for chunk in model:
+            for param in chunk.parameters():
+                group_name = getattr(param, "_lora_grad_sum_group", None)
+                if group_name is not None and param.requires_grad:
+                    marked.append((param, group_name))
+        _marked_lora_grad_params_cache[key] = marked
+    if not marked:
+        return
+    groups = {
+        "tp": (ps.get_tensor_model_parallel_group(), ps.get_tensor_model_parallel_world_size()),
+        "ep": (ps.get_expert_model_parallel_group(), ps.get_expert_model_parallel_world_size()),
+    }
+    for group_name in ("tp", "ep"):
+        group, size = groups[group_name]
+        if size <= 1:
+            continue
+        grads = []
+        for param, g_name in marked:
+            if g_name != group_name:
+                continue
+            grad = getattr(param, "main_grad", None)
+            if grad is None:
+                grad = param.grad
+            if grad is not None:
+                grads.append(grad)
+        for dt in {g.dtype for g in grads}:
+            gs = [g for g in grads if g.dtype == dt]
+            if len(gs) == 1:
+                dist.all_reduce(gs[0], op=dist.ReduceOp.SUM, group=group)
+                continue
+            flat = torch._utils._flatten_dense_tensors(gs)
+            dist.all_reduce(flat, op=dist.ReduceOp.SUM, group=group)
+            for g, red in zip(gs, torch._utils._unflatten_dense_tensors(flat, gs), strict=False):
+                g.copy_(red)
+
+
 def is_lora_model(model: Sequence[torch.nn.Module]) -> bool:
     """Check if model has LoRA layers applied."""
     for model_chunk in model:
@@ -365,7 +419,7 @@ def save_lora_checkpoint(
     1. **HF PEFT format** (``adapter_model.bin`` + ``adapter_config.json``) for
        external tool compatibility. Uses Megatron-Bridge's ``export_adapter_weights``
        which correctly handles fused QKV / gate-up weight splitting and TP gathering.
-    2. **Megatron-native format** (``adapter_megatron_tp{tp}_pp{pp}.pt``) for fast
+    2. **Megatron-native format** (``adapter_megatron_rank{global_rank}.pt``) for fast
        checkpoint resume without name/weight conversion. Each TP/PP rank saves its
        own shard with original parameter names.
 
@@ -388,61 +442,63 @@ def save_lora_checkpoint(
     tp_rank = parallel_state.tp.rank
     pp_rank = parallel_state.pp.rank
 
-    # Create directory on dp_rank=0, then synchronize
-    if is_dp_cp_rank_0:
-        save_path.mkdir(parents=True, exist_ok=True)
+    save_path.mkdir(parents=True, exist_ok=True)
     if dist.is_initialized():
         dist.barrier()
 
-    # ---- Megatron-native format (per TP/PP rank, fast resume) ----
-    if is_dp_cp_rank_0:
-        adapter_state: dict[str, torch.Tensor] = {}
-        for model_chunk in model:
-            for name, param in model_chunk.named_parameters():
-                if _is_adapter_param_name(name):
-                    adapter_state[name] = param.data.cpu()
+    adapter_state: dict[str, torch.Tensor] = {}
+    for model_chunk in model:
+        for name, param in model_chunk.named_parameters():
+            if _is_adapter_param_name(name):
+                adapter_state[name] = param.data.cpu()
 
-        native_path = save_path / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt"
-        torch.save(adapter_state, native_path)
-        logger.info(f"Saved {len(adapter_state)} adapter tensors (native) to {native_path}")
+    global_rank = dist.get_rank() if dist.is_initialized() else 0
+    native_path = save_path / f"adapter_megatron_rank{global_rank}.pt"
+    torch.save(adapter_state, native_path)
+    logger.info(f"Saved {len(adapter_state)} adapter tensors (native) to {native_path}")
 
     # ---- HF PEFT format (uses bridge for correct name/weight conversion) ----
     # Bridge export is collective: all TP ranks participate in the all-gather,
     # so every rank must call export_adapter_weights.
-    bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
-
-    lora_state_dict: dict[str, torch.Tensor] = {}
-    with megatron_bridge_utils.patch_megatron_model(model):
-        for hf_name, weight, _megatron_name in bridge.export_adapter_weights(
-            model,
-            cpu=True,
-            show_progress=False,
-        ):
-            lora_state_dict[hf_name] = weight
-
-    # Only one rank writes the HF PEFT files (bridge already gathered across TP/PP)
-    if is_dp_cp_rank_0 and tp_rank == 0 and pp_rank == 0:
-        torch.save(lora_state_dict, save_path / "adapter_model.bin")
-
-        target_modules_hf = (
-            convert_target_modules_to_hf(list(args.target_modules))
-            if args.target_modules
-            else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+    try:
+        bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
+
+        lora_state_dict: dict[str, torch.Tensor] = {}
+        with megatron_bridge_utils.patch_megatron_model(model):
+            for hf_name, weight, _megatron_name in bridge.export_adapter_weights(
+                model,
+                cpu=True,
+                show_progress=False,
+            ):
+                lora_state_dict[hf_name] = weight
+
+        if is_dp_cp_rank_0 and tp_rank == 0 and pp_rank == 0:
+            torch.save(lora_state_dict, save_path / "adapter_model.bin")
+
+            target_modules_hf = (
+                convert_target_modules_to_hf(list(args.target_modules))
+                if args.target_modules
+                else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+            )
+            config = {
+                "peft_type": "LORA",
+                "r": args.lora_rank,
+                "lora_alpha": args.lora_alpha,
+                "target_modules": target_modules_hf,
+                "lora_dropout": args.lora_dropout,
+                "bias": "none",
+                "task_type": "CAUSAL_LM",
+            }
+            with open(save_path / "adapter_config.json", "w") as f:
+                json.dump(config, f, indent=2)
+
+            os.sync()
+            logger.info(f"Saved HF PEFT adapter to {save_path} with {len(lora_state_dict)} tensors")
+    except Exception as hf_export_err:
+        logger.warning(
+            f"HF PEFT adapter export skipped ({hf_export_err}); the per-rank native "
+            f"shards + training state are sufficient for training resume."
         )
-        config = {
-            "peft_type": "LORA",
-            "r": args.lora_rank,
-            "lora_alpha": args.lora_alpha,
-            "target_modules": target_modules_hf,
-            "lora_dropout": args.lora_dropout,
-            "bias": "none",
-            "task_type": "CAUSAL_LM",
-        }
-        with open(save_path / "adapter_config.json", "w") as f:
-            json.dump(config, f, indent=2)
-
-        os.sync()
-        logger.info(f"Saved HF PEFT adapter to {save_path} with {len(lora_state_dict)} tensors")
 
     # ---- Training state (optimizer + scheduler) for resume ----
     if optimizer is not None:
@@ -500,7 +556,13 @@ def load_lora_adapter(
     pp_rank = get_parallel_state().pp.rank
 
     # ---- Try Megatron-native format first (fast, no conversion needed) ----
-    native_path = adapter_dir / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt"
+    global_rank = dist.get_rank() if dist.is_initialized() else 0
+    native_path = adapter_dir / f"adapter_megatron_rank{global_rank}.pt"
+    if not native_path.exists():
+        legacy = adapter_dir / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt"
+        if legacy.exists():
+            logger.warning(f"Using legacy tp/pp-named adapter shard {legacy}; only valid when EP<=TP")
+            native_path = legacy
     if native_path.exists():
         state_dict = torch.load(native_path, map_location="cpu", weights_only=True)
         loaded = 0
@@ -566,11 +628,14 @@ def _load_training_state(
 
 def build_lora_sync_config(args: Namespace) -> dict[str, Any]:
     """Build LoRA config dict for syncing weights to SGLang engines."""
-    target_modules_hf = (
-        target_modules_hf_for_sglang_rollout(args)
-        if args.target_modules
-        else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
-    )
+    if sglang_lora_target_all_sentinel(args):
+        target_modules_hf: Any = "all-linear"
+    else:
+        target_modules_hf = (
+            target_modules_hf_for_sglang_rollout(args)
+            if args.target_modules
+            else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+        )
     return {
         "peft_type": "LORA",
         "r": args.lora_rank,

miles/backends/megatron_utils/megatron_to_hf/__init__.py

@@ -2,6 +2,7 @@
 from .deepseekv4 import convert_deepseekv4_to_hf
 from .glm4 import convert_glm4_to_hf
 from .glm4moe import convert_glm4moe_to_hf
+from .inkling import convert_inkling_to_hf
 from .kimi_vl import convert_kimi_k25_to_hf, convert_kimivl_to_hf
 from .llama import convert_llama_to_hf
 from .mimo import convert_mimo_to_hf
@@ -52,6 +53,8 @@ def _convert_to_hf_core(args, model_name, name, param):
         converted_named_tensors = convert_qwen2_to_hf(args, name, param)
     elif "deepseekv4" in model_name:
         converted_named_tensors = convert_deepseekv4_to_hf(args, name, param)
+    elif "inkling" in model_name:
+        converted_named_tensors = convert_inkling_to_hf(args, name, param)
     elif "llama" in model_name:
         converted_named_tensors = convert_llama_to_hf(args, name, param)
     elif "mimo" in model_name: