feat(unsloth): support Nemotron-H hybrid MoE fine-tuning on DGX Spark (#338)

* feat(unsloth): support Nemotron-H hybrid Mamba MoE fine-tuning

Enable fine-tuning of hybrid Mamba-2 + MoE models (NVIDIA Nemotron-H,
e.g. Nemotron-3-Nano-30B-A3B) on the Unsloth customization backend, plus
two supporting fixes uncovered bringing the backend up on a unified-memory
GB10 (DGX Spark, arm64/Blackwell).

- Add mamba-ssm + causal-conv1d to the nmp-unsloth-training image, installed
  from the same prebuilt cu13.1.1/cp312 wheels as nmp-automodel-base via the
  causal-conv1d-wheel / mamba-ssm-wheel bake contexts (versions single-sourced,
  no recompile). Nemotron-H remote code imports mamba_ssm at load time.
- Add configurable model.device_map (default {"": 0}). Without an explicit
  single-device map, accelerate's device_map="auto" sizes the GPU budget from
  CUDA free memory, which on unified-memory parts reports a small dynamic
  carve-out and spills layers to CPU -- aborting 4-bit loads.
- Fix file_io download crash: get_percentage now clamps instead of raising when
  the live file count exceeds the pre-listing total (off-by-one on sources with
  nested directories), so a cosmetic progress value can't fail a multi-GB
  download mid-transfer.
- Register cpu/gpu + gpu/gpu docker execution profiles in the local runner
  config so the customization backends run out-of-the-box on a local platform.
- Docs: document model.device_map and the 16-bit-LoRA path for hybrid MoE models
  in the customizer hyperparameters reference; refresh the unsloth image README
  build steps (shared wheels / USE_LOCAL_WHEELS, accurate step list).

Validated end-to-end on DGX Spark: Nemotron-3-Nano-30B-A3B loads and trains in
16-bit LoRA (load_in_4bit=false).

Signed-off-by: Albert Cui <albcui@nvidia.com>

* add docker backend to models controller

Signed-off-by: Albert Cui <albcui@nvidia.com>

---------

Signed-off-by: Albert Cui <albcui@nvidia.com>

Author: albcui

docker-bake.hcl

@@ -796,7 +796,9 @@ target "nmp-unsloth-training" {
   dockerfile = "docker/Dockerfile.nmp-unsloth-training"
   target     = "runtime"
   contexts = {
-    platform-workspace = "target:unsloth-platform-workspace"
+    platform-workspace        = "target:unsloth-platform-workspace"
+    causal-conv1d-wheel-image = causal_conv1d_wheel_context()
+    mamba-ssm-wheel-image     = mamba_ssm_wheel_context()
   }
   cache-to   = maybe_registry_cache_to("nmp-unsloth-training")
   cache-from = maybe_registry_cache_from("nmp-unsloth-training")

docker/Dockerfile.nmp-unsloth-training

@@ -15,10 +15,15 @@
 #   1b. bitsandbytes — compiled from source against NGC CUDA 13.1 (same pattern
 #       as docker/Dockerfile.nmp-automodel-base). PyPI wheels
 #       only ship through cuda130; source build replaces the wheel from step 1.
-#   1c. flash-attn — optional for unsloth and not installed by step 1. Without
-#       it Unsloth falls back when xformers is also missing (common on newer CUDA
-#       stacks), logging "FA2 = False / Xformers = None". Installed immediately
-#       after unsloth so pip does not re-resolve the HF stack.
+#   1c. mamba-ssm + causal-conv1d — installed from the prebuilt cu13.1.1 / cp312
+#       wheels shared with docker/Dockerfile.nmp-automodel-base (the
+#       causal-conv1d-wheel / mamba-ssm-wheel bake contexts). Required by hybrid
+#       Mamba/SSM models (e.g. NVIDIA Nemotron-H *-A3B) whose remote code imports
+#       mamba_ssm at load time.
+#   1d. flash-attn — optional for unsloth and currently NOT installed (see the
+#       commented TODO below). Without it Unsloth falls back when xformers is
+#       also missing (common on newer CUDA stacks), logging "FA2 = False /
+#       Xformers = None".
 #   2. Editable installs of the platform glue (nemo-platform SDK, plugin,
 #      nmp-common, nmp-unsloth) from the in-repo workspace slice.
 #
@@ -29,6 +34,14 @@
 # Override at build time: --set nmp-unsloth-training.args.PYTORCH_BASE=...
 ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.02-py3
 
+# Prebuilt CUDA-extension wheels (mamba-ssm + causal-conv1d), shared with
+# docker/Dockerfile.nmp-automodel-base. The bake `nmp-unsloth-training` target
+# resolves these contexts to the causal-conv1d-wheel / mamba-ssm-wheel images
+# (see docker-bake.hcl + docker/base/Dockerfile.mamba-wheel). Local builds need
+# `USE_LOCAL_WHEELS=1` (build the wheel targets) or a reachable WHEELS_REGISTRY.
+FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
+FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src
+
 FROM ${PYTORCH_BASE} AS base
 
 WORKDIR /opt
@@ -105,7 +118,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip uninstall --python ${VIRTUAL_ENV}/bin/python scikit-build-core && \
     rm -rf /tmp/bitsandbytes
 
-# TODO: Step 1c: Flash Attention 2 — compiled from source against the NGC 26.02 torch.
+# Step 1c: mamba-ssm + causal-conv1d — required by hybrid Mamba/SSM models such
+# as NVIDIA Nemotron-H (e.g. *-Nano-*-A3B), whose remote modeling code imports
+# mamba_ssm at load time. Installed from the SAME prebuilt cu13.1.1 / cp312
+# wheels as docker/Dockerfile.nmp-automodel-base — built once by the
+# causal-conv1d-wheel / mamba-ssm-wheel bake targets (docker/base/Dockerfile.mamba-wheel)
+# and shared across both images, so the versions are single-sourced and nothing
+# recompiles here. The wheels are built from this same NGC 26.02 base (torch 2.11
+# + CUDA 13.1), so they're ABI-matched. --no-deps so the HF stack isn't
+# re-resolved (triton is already provided by the NGC base).
+RUN --mount=from=causal-conv1d-wheel-src,target=/tmp/causal-conv1d-wheel-src,readonly \
+    --mount=from=mamba-ssm-wheel-src,target=/tmp/mamba-ssm-wheel-src,readonly \
+    uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --no-deps \
+    /tmp/causal-conv1d-wheel-src/wheels/cu13.1.1/causal_conv1d-*cp312*.whl \
+    /tmp/mamba-ssm-wheel-src/wheels/cu13.1.1/mamba_ssm-2.3.0-cp312*.whl
+
+# TODO: Step 1d: Flash Attention 2 — compiled from source against the NGC 26.02 torch.
 # /usr/local/cuda symlinks to an older toolkit; use /usr/local/cuda-13.1 instead.
 # Cap parallel nvcc/ninja work — default uses all CPUs and OOMs typical build hosts.
 # Put flash attention back in when we have a working wheel in a separate image.

docker/unsloth/README.md

@@ -46,21 +46,31 @@ docker buildx bake \
 # Result: `${IMAGE_REGISTRY}/nmp-unsloth-training:${BAKE_TAG}` in the registry.
 ```
 
-The build pulls the NGC PyTorch base, then runs unsloth's canonical install
-in two steps:
+> **Prebuilt wheels (local builds).** The image installs `mamba-ssm` +
+> `causal-conv1d` from the shared `causal-conv1d-wheel` / `mamba-ssm-wheel`
+> bake contexts (same wheels as `nmp-automodel-base`). A **local** bake must
+> build those wheels first — prepend `USE_LOCAL_WHEELS=1` (or point
+> `WHEELS_REGISTRY`/`WHEELS_TAG` at prebuilt wheels), otherwise bake tries to
+> pull `${WHEELS_REGISTRY}/causal-conv1d-wheel:${WHEELS_TAG}` and fails. To
+> avoid recompiling on every image build, build the wheels once with
+> `USE_LOCAL_WHEELS=1 docker buildx bake -f docker-bake.hcl nmp-automodel-gpu-wheels`.
+
+The build pulls the NGC PyTorch base, then:
 
 1. `uv pip install unsloth --torch-backend=auto transformers==4.57.6 huggingface-hub==0.36.2` with
    `preserve_base_torch.txt` overrides so the NGC base's PyTorch + CUDA are not
    replaced. Unsloth's resolver still pulls `unsloth_zoo`, trl, peft,
    accelerate, datasets, bitsandbytes, and xformers. **transformers is pinned
    explicitly** to `4.57.6` (override at build time via
    `--build-arg TRANSFORMERS_VERSION=...`).
-1b. Flash Attention 2 (Dockerfile step 1b) — source build with
-    `--no-build-isolation` against the NGC base torch (cached Docker layer).
-    Parallelism is capped via `MAX_JOBS` (default `2`, override with bake arg
-    `FLASH_ATTN_MAX_JOBS`) and `NVCC_THREADS=1` to avoid OOM during nvcc.
-    Unsloth does not depend on `flash-attn`; without it you may see
-    `FA2 = False` and `Xformers = None` on newer CUDA stacks.
+1b. bitsandbytes — compiled from source against the NGC CUDA 13.1 toolkit
+    (PyPI wheels only ship through cuda130), replacing the wheel from step 1.
+1c. mamba-ssm + causal-conv1d — installed from the prebuilt `cu13.1.1` / `cp312`
+    wheels shared with `nmp-automodel-base` (see the **Prebuilt wheels** note
+    above). Required by hybrid Mamba/SSM models (e.g. NVIDIA Nemotron-H `*-A3B`).
+1d. Flash Attention 2 — **not currently installed** (commented TODO in the
+    Dockerfile). Unsloth does not depend on it; without it you may see
+    `FA2 = False` / `Xformers = None` on newer CUDA stacks.
 2. Editable install of the platform glue: `nemo-platform-sdk`,
    `nemo-platform-plugin`, `nmp-common`, `nmp-unsloth`.
 
@@ -140,12 +150,19 @@ docker buildx bake \
 ```bash
 export IMAGE_REGISTRY="my-registry/nemo-platform-dev"
 export BAKE_TAG="$(git rev-parse --short HEAD)"
+# The image needs the mamba-ssm + causal-conv1d wheel images as build contexts.
+# A direct `docker buildx build` can't resolve the bake `*_wheel_context()`
+# functions, so prefer `docker buildx bake nmp-unsloth-training` (with
+# USE_LOCAL_WHEELS=1) which wires them automatically. If you must use
+# `docker buildx build`, pass both contexts explicitly:
 docker buildx build \
   -f docker/Dockerfile.nmp-unsloth-training \
   --output type=docker,dest=/tmp/nmp-unsloth-training.tar \
   --target runtime \
   -t "${IMAGE_REGISTRY}/nmp-unsloth-training:${BAKE_TAG}" \
   --build-context "platform-workspace=path-to-platform-workspace" \
+  --build-context "causal-conv1d-wheel-image=docker-image://${WHEELS_REGISTRY}/causal-conv1d-wheel:${WHEELS_TAG}" \
+  --build-context "mamba-ssm-wheel-image=docker-image://${WHEELS_REGISTRY}/mamba-ssm-wheel:${WHEELS_TAG}" \
   .
 
 scp /tmp/nmp-unsloth-training.tar gpu-pod:/tmp/

packages/nmp_platform/config/local.yaml

@@ -71,17 +71,20 @@ jobs:
         ttl_seconds_before_active: 60
         ttl_seconds_active: 3600
         ttl_seconds_after_finished: 300
-    # Uncomment for using customizer
-    # - provider: cpu
-    #   profile: gpu
-    #   backend: docker
-    #   config:
-    #     launcher_tool_path: ./services/core/jobs/jobs-launcher/jobs-launcher
-    # - provider: gpu
-    #   profile: gpu
-    #   backend: docker
-    #   config:
-    #     launcher_tool_path: ./services/core/jobs/jobs-launcher/jobs-launcher
+    # Customizer (automodel / unsloth) GPU jobs. The unsloth backend stamps the
+    # profile name "gpu" onto all 4 steps, so both cpu/gpu and gpu/gpu must
+    # exist. `profile` is a name, not a resource claim — `provider` governs
+    # CPU vs GPU, so provider: cpu + profile: gpu is correct (not a mismatch).
+    - provider: cpu
+      profile: gpu
+      backend: docker
+      config:
+        launcher_tool_path: ./services/core/jobs/jobs-launcher/jobs-launcher
+    - provider: gpu
+      profile: gpu
+      backend: docker
+      config:
+        launcher_tool_path: ./services/core/jobs/jobs-launcher/jobs-launcher
 
   # Local path to the jobs-launcher binary used by the Docker job backend
   executor_defaults:
@@ -118,6 +121,12 @@ models:
     backends:
       docker:
         enabled: true
+        # LoRA adapter-download sidecar for deployments with lora_enabled. The
+        # docker backend has no separate args field (unlike k8s) and keeps the
+        # image ENTRYPOINT — nmp-automodel-tasks is `/opt/venv/bin/python`, so
+        # the module goes in lora_sidecar_command (runs `python -m <module>`).
+        lora_sidecar_image_name: nmp-automodel-tasks
+        lora_sidecar_command: ["-m", "nmp.core.models.sidecars.adapters.main"]
 
 # Inference gateway configuration - uses default values
 inference_gateway: {}

packages/nmp_platform_runner/src/nmp/platform_runner/config/local.yaml

@@ -28,17 +28,20 @@ jobs:
         ttl_seconds_before_active: 60
         ttl_seconds_active: 3600
         ttl_seconds_after_finished: 300
-    # Uncomment for using customizer
-    # - provider: cpu
-    #   profile: gpu
-    #   backend: docker
-    #   config:
-    #     launcher_tool_path: ./services/core/jobs/jobs-launcher/jobs-launcher
-    # - provider: gpu
-    #   profile: gpu
-    #   backend: docker
-    #   config:
-    #     launcher_tool_path: ./services/core/jobs/jobs-launcher/jobs-launcher
+    # Customizer (automodel / unsloth) GPU jobs. The unsloth backend stamps the
+    # profile name "gpu" onto all 4 steps, so both cpu/gpu and gpu/gpu must
+    # exist. `profile` is a name, not a resource claim — `provider` governs
+    # CPU vs GPU, so provider: cpu + profile: gpu is correct (not a mismatch).
+    - provider: cpu
+      profile: gpu
+      backend: docker
+      config:
+        launcher_tool_path: ./services/core/jobs/jobs-launcher/jobs-launcher
+    - provider: gpu
+      profile: gpu
+      backend: docker
+      config:
+        launcher_tool_path: ./services/core/jobs/jobs-launcher/jobs-launcher
 
   executor_defaults:
     docker:
@@ -57,6 +60,15 @@ models:
   controller:
     interval_seconds: 5
     model_deployment_garbage_collection_ttl_seconds: 30
+    backends:
+      docker:
+        enabled: true
+        # LoRA adapter-download sidecar for deployments with lora_enabled. The
+        # docker backend has no separate args field (unlike k8s) and keeps the
+        # image ENTRYPOINT — nmp-automodel-tasks is `/opt/venv/bin/python`, so
+        # the module goes in lora_sidecar_command (runs `python -m <module>`).
+        lora_sidecar_image_name: nmp-automodel-tasks
+        lora_sidecar_command: ["-m", "nmp.core.models.sidecars.adapters.main"]
 
 inference_gateway: {}
 

plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/hyperparameters.md

@@ -462,10 +462,13 @@ Full template (every section, defaults inline):
 | `load_in_4bit` | `true` | bitsandbytes 4-bit. Mutex with `load_in_8bit`. Default for Unsloth's headline path; required to fit larger models on small GPUs. |
 | `load_in_8bit` | `false` | bitsandbytes 8-bit. Mutex with `load_in_4bit`. |
 | `dtype` | `"auto"` | One of `"auto"`, `"bfloat16"`, `"float16"`, `"float32"`. |
-| `trust_remote_code` | `false` | HF `trust_remote_code` flag for custom model code. |
+| `trust_remote_code` | `false` | HF `trust_remote_code` flag for custom model code (required by some hybrid Mamba/MoE models, e.g. Nemotron-H). |
+| `device_map` | `null` | Placement for `FastLanguageModel.from_pretrained`. `null` pins the whole model to the single visible GPU (`{"": 0}`) — the right default for this single-GPU backend. Leave unset unless experimenting; `"auto"`/`"balanced"`/`"sequential"` can spill layers to CPU on unified-memory hosts (GB10 / DGX Spark) and abort 4-bit loads. |
 
 **Mutex:** `load_in_4bit` xor `load_in_8bit`. Both quantization flags are also **incompatible with `training.finetuning_type: "all_weights"`** — full SFT must use a non-quantized base.
 
+> **Hybrid Mamba/MoE models (e.g. NVIDIA Nemotron-H `*-A3B`):** load in **16-bit** (`load_in_4bit: false`, `load_in_8bit: false`) — Unsloth's supported path for these. The 4-bit (bitsandbytes) path can hit a dtype mismatch inside the model's MoE expert accumulation. Keep `device_map` unset (single-GPU default) and set `trust_remote_code: true`.
+
 ### `dataset`
 
 See `references/dataset-formats.md` § Unsloth for row-shape rules.

services/unsloth/src/nmp/unsloth/schemas.py

@@ -61,6 +61,19 @@ class ModelLoadSpec(BaseModel):
     load_in_8bit: bool = False
     dtype: Literal["auto", "bfloat16", "float16", "float32"] = "auto"
     trust_remote_code: bool = False
+    device_map: str | int | dict[str, int] | None = Field(
+        default=None,
+        description=(
+            "Device placement forwarded to FastLanguageModel.from_pretrained. "
+            "Omit (null) to pin the whole model to the single visible GPU "
+            "({'': 0}) — the right default for this single-GPU backend, and it "
+            "avoids accelerate's auto-placement under-sizing GPU memory on "
+            "unified-memory parts (e.g. GB10 / DGX Spark), which otherwise "
+            "spills layers to CPU and aborts 4-bit loads. Set 'auto', "
+            "'balanced', 'sequential', a device index, or a custom map for "
+            "multi-device experiments."
+        ),
+    )
 
 
 class LoRAParams(BaseModel):

services/unsloth/src/nmp/unsloth/tasks/file_io/callbacks.py

@@ -20,22 +20,22 @@
 
 
 def get_percentage(current: int, total: int) -> int:
-    """Get integer percentage 0-100.
-
-    Raises:
-        ValueError: If current > total or either value is negative.
+    """Get integer percentage 0-100, clamped to the valid range.
+
+    Progress accounting must never abort the underlying transfer. Inputs
+    can fall outside ``[0, total]`` for benign reasons — most commonly when
+    the pre-transfer file listing under-counts a source that contains nested
+    directories, so the live ``current`` count exceeds ``total`` by the
+    number of nested files. Clamp rather than raise so a cosmetic progress
+    number can't fail a multi-GB download.
     """
-    if current > total:
-        raise ValueError(
-            f"Unexpected value of the current and total values: current={current} cannot be greater than total={total}",
-        )
-    if total < 0:
-        raise ValueError(f"Unexpected negative value of the total value: total={total}, current={current}")
-    if current < 0:
-        raise ValueError(f"Unexpected negative value of the current value: current={current}, total={total}")
-
-    if total == 0:
+    if total <= 0:
         return 0
+    if current > total or current < 0:
+        # Benign (see docstring) but worth a breadcrumb now that we no longer
+        # raise — the old hard error is what previously surfaced count drift.
+        logger.debug("get_percentage clamping out-of-range progress: current=%s total=%s", current, total)
+    current = max(0, min(current, total))
     return int((current / total) * 100)
 
 

services/unsloth/src/nmp/unsloth/tasks/training/backends/unsloth_sft.py

@@ -74,6 +74,7 @@ def build_model_load_kwargs(spec: UnslothJobOutput, resolved_model: str) -> dict
         "load_in_8bit": spec.model.load_in_8bit,
         "full_finetuning": spec.training.finetuning_type == "all_weights",
         "trust_remote_code": spec.model.trust_remote_code,
+        "device_map": spec.model.device_map if spec.model.device_map is not None else {"": 0},
     }