feat(RL): NeMo-RL × Megatron ModelExpress v2 weight-refit (trainer path)

docs/index.md

@@ -332,6 +332,7 @@ design-docs/training-backends.md
 design-docs/sequence-packing-and-dynamic-batching.md
 design-docs/env-vars.md
 design-docs/nemo-gym-integration.md
+design-docs/dynamo-integration.md
 ```
 
 ```{toctree}

examples/nemo_gym/grpo_workplace_assistant_dynamo_mx_gp.yaml

@@ -0,0 +1,115 @@
+# Workplace-assistant + Dynamo MX refit through a hierarchical
+# GlobalPlanner/GlobalRouter topology.
+#
+# Paired infra:
+#   infra/nrl_k8s/examples/grpo_workplace_assistant_dynamo_mx_gp.gb300.infra.yaml
+# Paired DGDs:
+#   infra/nrl_k8s/examples_dgd/qwen3_4b_thinking_gb300_mx_gp_{pool0,pool1,ctrl}.yaml
+#
+# The DGD name and refit worker namespaces are injected by the infra entrypoint.
+
+defaults: grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
+
+checkpointing:
+  enabled: false
+
+grpo:
+  num_prompts_per_step: 2
+  num_generations_per_prompt: 2
+  max_rollout_turns: 1
+  max_num_steps: 2
+  val_period: 999999
+  val_at_start: false
+  val_at_end: false
+
+policy:
+  model_name: "Qwen/Qwen3-4B-Thinking-2507"
+  tokenizer:
+    name: ${policy.model_name}
+    chat_template_kwargs: null
+  train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}
+  train_micro_batch_size: 1
+  generation_batch_size: ${policy.train_global_batch_size}
+  logprob_batch_size: 1
+  max_total_sequence_length: 16384
+  precision: "bfloat16"
+
+  megatron_cfg:
+    enabled: false
+
+  dtensor_cfg:
+    _v2: false
+    enabled: true
+    cpu_offload: false
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  optimizer:
+    name: "torch.optim.AdamW"
+    kwargs:
+      lr: 1.0e-6
+      weight_decay: 0.0
+      betas: [0.9, 0.999]
+      eps: 1e-8
+  scheduler:
+    - name: "torch.optim.lr_scheduler.ConstantLR"
+      kwargs:
+        factor: 1.0
+        total_iters: 10000000000
+    - milestones: []
+
+  generation:
+    backend: "dynamo"
+    max_new_tokens: ${policy.max_total_sequence_length}
+    temperature: 1.0
+    top_p: 1.0
+    top_k: null
+    stop_token_ids: null
+    stop_strings: null
+    vllm_cfg:
+      async_engine: true
+      precision: ${policy.precision}
+      kv_cache_dtype: "auto"
+      tensor_parallel_size: 1
+      pipeline_parallel_size: 1
+      gpu_memory_utilization: 0.75
+      max_model_len: ${policy.max_total_sequence_length}
+      enforce_eager: false
+      expose_http_server: true
+    vllm_kwargs: {}
+    dynamo_cfg: {}
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 1
+        num_nodes: 1
+
+env:
+  nemo_gym:
+    policy_model:
+      responses_api_models:
+        vllm_model:
+          uses_reasoning_parser: true
+          extra_body:
+            chat_template_kwargs:
+              enable_thinking: true
+
+cluster:
+  gpus_per_node: 1
+  num_nodes: 1
+  weight_sync:
+    method: "mx"
+    mx_config:
+      enabled: true
+      mx_server_url: "modelexpress-server.default.svc.cluster.local:8001"
+      timeout_seconds: 300.0
+      same_rank_only: true
+      tree_scale_out: true
+      moe_expert_filter: false
+      nic_pin: "auto"
+
+logger:
+  log_dir: "logs/grpo-workplace-assistant-dynamo-mx-gp"

infra/helm/helmfile.yaml

@@ -61,3 +61,17 @@ releases:
   version: 0.11.1
   wait: true
 
+#
+# Dynamo Platform: dynamo-operator (reconciles DynamoGraphDeployment CRs into
+# pods), bundled etcd + NATS for service discovery / event plane. Required when
+# any nrl-k8s recipe sets `infra.dynamo.<key>`.
+#
+- name: dynamo-platform
+  namespace: dynamo-platform
+  createNamespace: true
+  chart: oci://ghcr.io/ai-dynamo/dynamo-platform
+  version: 1.2.0
+  wait: true
+  values:
+  - values/dynamo-platform.yaml
+

infra/helm/values/dynamo-platform.yaml

@@ -0,0 +1,24 @@
+# Dynamo Platform values for the nrl-k8s helmfile.
+#
+# This integrates with the kai-scheduler we already install (v0.14.0) and uses
+# the platform chart's bundled etcd + NATS so the operator has the service
+# discovery and event plane it expects out of the box.
+#
+# To enable Grove (multi-node DGDs), flip both `install` and `enabled` under
+# `global.grove`. Off here because Phase 2 only ships single-node DGDs.
+
+global:
+  etcd:
+    install: true
+  kai-scheduler:
+    install: false   # already installed by the kai-scheduler release in helmfile.yaml
+    enabled: true    # the operator should still inject schedulerName + queue labels into DGD pods
+  grove:
+    install: false
+    enabled: false
+
+dynamo-operator:
+  enabled: true
+  upgradeCRD: true   # apply CRDs via the pre-install/pre-upgrade hook Job
+  nats:
+    enabled: true

infra/nrl_k8s/dynamo_mx/Dockerfile

@@ -0,0 +1,60 @@
+# syntax=docker/dockerfile:1.6
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Overlay image: base Dynamo vLLM worker + ModelExpress v2 Python client + NIXL.
+#
+# The base image (BASE_IMAGE arg) is the same arm64 + Blackwell build we use
+# for the non-MX dynamo path
+# (e.g. jwillthomson/dynamo-arm-rl-tokenize-endpoint-91cdc71:latest). This
+# overlay adds:
+#
+#   * The `modelexpress` Python client from `kavink/nemo_rl_moe` on
+#     ai-dynamo/modelexpress (head e0c9447939f0420f6815464b09455724e57d7553
+#     as of 2026-06-23) — contains
+#     MxV2RefitReceiver, the shape registry codecs, and the picker.
+#   * NIXL userspace + UCX bindings, required for RDMA transfers.
+#
+# The Dynamo-side refit Python code (dynamo.vllm.mx_refit.*) is expected to be
+# part of the base image, built from the Dynamo integration branch. This
+# Dockerfile only pins the ModelExpress client/NIXL layer; it does not copy a
+# runtime source overlay from the NeMo-RL checkout.
+#
+# Build:
+#   docker buildx build --platform linux/arm64 \
+#     --build-arg BASE_IMAGE=jwillthomson/dynamo-arm-rl-tokenize-endpoint-91cdc71:latest \
+#     --build-arg MX_REF=e0c9447939f0420f6815464b09455724e57d7553 \
+#     -t jwillthomson/dynamo-arm-rl-tokenize-endpoint-91cdc71-mx:e0c9447 \
+#     -f infra/nrl_k8s/dynamo_mx/Dockerfile .
+#
+# Push:
+#   docker push jwillthomson/dynamo-arm-rl-tokenize-endpoint-91cdc71-mx:e0c9447
+
+ARG BASE_IMAGE=jwillthomson/dynamo-arm-rl-tokenize-endpoint-91cdc71:latest
+FROM ${BASE_IMAGE}
+
+# `kavink/nemo_rl_moe` is the v2-refit branch. Pin to a SHA so rebuilds are
+# reproducible.
+ARG MX_REF=e0c9447939f0420f6815464b09455724e57d7553
+
+# NIXL: GitHub release of nixl + ucx. The dynamo image typically already
+# ships UCX; this just ensures the Python bindings match.
+# TODO: confirm whether the base image already has nixl-python. If yes,
+# remove the pip install below.
+
+USER root
+
+# System deps for NIXL: librdmacm, libibverbs already in the base via mlnx_ofed;
+# we just need to pip-install the bindings + the modelexpress client.
+RUN python3 -m pip install --no-cache-dir \
+      "git+https://github.com/ai-dynamo/modelexpress.git@${MX_REF}#subdirectory=modelexpress_client/python" \
+ && python3 -c "from modelexpress import MxV2RefitReceiver, MxV2TrainingPublisher; print('modelexpress imports OK')"
+
+# Smoke: import the dynamo refit extension + check it can be resolved as a
+# worker_extension_cls.
+RUN python3 -c "\
+from dynamo.vllm.mx_refit.extension import MxRefitWorkerExtension; \
+assert hasattr(MxRefitWorkerExtension, 'update_weights_via_mx'); \
+print('MxRefitWorkerExtension import OK')"
+
+# Default to the same entrypoint as the base — DGD manifests override this.