feat(distillation): add experimental Nano3 TransferQueue QAOPD

Signed-off-by: Meng Xin <mxin@nvidia.com>

Author: mxinO

examples/configs/distillation_math.yaml

@@ -17,6 +17,28 @@ loss_fn:
     kl_type: "mixed" # forward, reverse, mixed
     mixed_kl_weight: 0.5 # when kl_type is "mixed", this is the weight of the forward KL
     zero_outside_topk: false # zero out the teacher logits outside the top k when calculate forward KL loss
+    sft_weight: 0.0
+    entropy_diagnostics:
+        enabled: false
+        teacher_entropy_low: 0.3
+        teacher_entropy_high: 0.7
+        log_position_buckets: true
+    entropy_aware:
+        enabled: false
+        teacher_entropy_threshold: 0.8
+        weighting_mode: "threshold" # threshold, linear_ramp
+        ramp_start: 0.3
+        ramp_width: 0.6
+        forward_kl_weight: 1.0
+        high_entropy_reverse_kl_weight: 1.0
+        high_entropy_reverse_kl_threshold: 0.7
+        low_entropy_sharpening_enabled: false
+        low_entropy_sharpening_weight: 0.0
+        low_entropy_sharpening_temperature: 0.7
+        low_entropy_sharpening_support_ratio_threshold: 1.0
+        low_entropy_sharpening_requires_support_ratio_above_threshold: true
+        require_sft_weight_zero: true
+        require_zero_outside_topk_false: true
 
 checkpointing:
     enabled: true

examples/modelopt/nano3_qaopd_tq/README.md

@@ -0,0 +1,127 @@
+# Nano3 QAOPD with TransferQueue
+
+This experimental example trains an NVFP4 Nemotron 3 Nano student against the
+matching BF16 teacher. Student generation produces NeMo-Gym trajectories, the
+teacher writes top-k logits directly to TransferQueue, and the Megatron policy
+worker consumes the queued fields during training.
+
+The supplied recipe preserves the development topology:
+
+- 8 nodes with 4 GPUs per node
+- Megatron TP=4, CP=4, EP=8
+- vLLM TP=4 and EP=4
+- 128 prompts with 4 generations per prompt
+- global training batch size 512
+- teacher top-k 300, BF16 logits, and int32 indices
+- `mooncake_cpu` storage with microbatch-streamed top-k writeback
+
+Treat these as starting values. Memory capacity and parallelism must be checked
+again when the GPU type, sequence length, or batch shape changes.
+
+## Build an Enroot image
+
+Enroot imports OCI/Docker images; it does not build a Dockerfile itself. Build
+the NeMo RL release image first, then convert it to a squashfs image.
+
+Initialize the repository and build on the same CPU architecture as the target
+GPU cluster whenever possible:
+
+```bash
+git submodule update --init --recursive
+
+docker buildx build \
+  --platform linux/arm64 \
+  --target release \
+  --build-context nemo-rl=. \
+  --file docker/Dockerfile \
+  --tag nemo-rl:nano3-tq \
+  --load \
+  .
+
+enroot import \
+  --output nemo-rl-nano3-tq.sqsh \
+  dockerd://nemo-rl:nano3-tq
+```
+
+Use `linux/amd64` for an x86_64 target. Cross-architecture builds require a
+Buildx builder with binfmt/QEMU support and are slower and less reliable for
+CUDA extension builds. A native build on the target architecture is preferred.
+
+If Docker and Enroot are on different machines, push to a registry and import
+the image there. The `#` separates the registry from the image path in Enroot's
+Docker URI syntax:
+
+```bash
+docker buildx build \
+  --platform linux/arm64 \
+  --target release \
+  --build-context nemo-rl=. \
+  --file docker/Dockerfile \
+  --tag registry.example.com/team/nemo-rl:nano3-tq \
+  --push \
+  .
+
+enroot import \
+  --arch arm64 \
+  --output nemo-rl-nano3-tq.sqsh \
+  docker://registry.example.com#team/nemo-rl:nano3-tq
+```
+
+The resulting `.sqsh` file can be passed directly to Pyxis through
+`--container-image` or the repository's `ray.sub` launcher.
+
+Do not put registry credentials, Hugging Face tokens, W&B keys, or other
+secrets in the image, scripts, Git configuration, or recipe. Authenticate with
+the registry before import and inject runtime credentials through your
+cluster's secret-management mechanism.
+
+## Prepare data
+
+The launcher expects two files on storage shared by every Ray node:
+
+1. NeMo-Gym JSONL training data accepted by `NemoGymDataset`.
+2. JSONL calibration data used to initialize the NVFP4 quantizers.
+
+The model may be a shared local directory or the public Hugging Face model ID
+used by the recipe. For gated models, configure Hugging Face authentication in
+the runtime environment without adding credentials to this repository.
+
+## Run on an existing Ray cluster
+
+From the repository root inside the container:
+
+```bash
+bash examples/modelopt/nano3_qaopd_tq/run_training.sh \
+  /shared/data/train.jsonl \
+  /shared/data/quant-calibration.jsonl \
+  /shared/results/nano3-qaopd-tq \
+  nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+```
+
+W&B is disabled in the example recipe. Checkpoints, JSON logs, and TensorBoard
+files are written below the supplied output directory.
+
+## Submit through Slurm, Pyxis, and Ray
+
+The submission wrapper takes all cluster-specific values as arguments. No
+account, partition, filesystem root, or credential is stored in the script.
+
+```bash
+bash examples/modelopt/nano3_qaopd_tq/submit_slurm.sh \
+  /shared/containers/nemo-rl-nano3-tq.sqsh \
+  SLURM_ACCOUNT \
+  SLURM_PARTITION \
+  /shared \
+  /shared/models/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
+  /shared/data/train.jsonl \
+  /shared/data/quant-calibration.jsonl \
+  /shared/results/nano3-qaopd-tq
+```
+
+The wrapper mounts only the specified shared root and the source checkout. The
+underlying `ray.sub` invocation uses exclusive nodes and disables automatic
+home-directory mounting inside the container.
+
+For a smoke test, copy the recipe, reduce only the number of steps and prompt
+count, and keep sequence length, model parallelism, quantization, top-k, and
+TQ dtype settings unchanged when validating memory behavior.

examples/modelopt/nano3_qaopd_tq/run_training.sh

@@ -0,0 +1,80 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+usage() {
+    cat <<'EOF'
+Usage:
+  run_training.sh TRAIN_DATA QUANT_CALIB_DATA OUTPUT_DIR [MODEL]
+
+Arguments:
+  TRAIN_DATA       NeMo-Gym JSONL training data on shared storage.
+  QUANT_CALIB_DATA JSONL calibration data on shared storage.
+  OUTPUT_DIR       Shared directory for checkpoints and logs.
+  MODEL            Local model path or Hugging Face model ID.
+
+Optional environment:
+  NUM_NODES        Ray worker nodes (default: 8).
+  GPUS_PER_NODE    GPUs per node (default: 4).
+  MAX_NUM_STEPS    Training steps (default: 3000).
+EOF
+}
+
+if (( $# < 3 || $# > 4 )); then
+    usage >&2
+    exit 2
+fi
+
+TRAIN_DATA=$1
+QUANT_CALIB_DATA=$2
+OUTPUT_DIR=$3
+MODEL=${4:-nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16}
+NUM_NODES=${NUM_NODES:-8}
+GPUS_PER_NODE=${GPUS_PER_NODE:-4}
+MAX_NUM_STEPS=${MAX_NUM_STEPS:-3000}
+
+if [[ ! -f "$TRAIN_DATA" ]]; then
+    echo "Training data does not exist: $TRAIN_DATA" >&2
+    exit 1
+fi
+if [[ ! -f "$QUANT_CALIB_DATA" ]]; then
+    echo "Calibration data does not exist: $QUANT_CALIB_DATA" >&2
+    exit 1
+fi
+
+mkdir -p "$OUTPUT_DIR"
+
+export DISABLE_MODELOPT_LAYER_SPEC=1
+export ENABLE_BRIDGE_QUANT_MAPPING=1
+export NRL_TQ_MOONCAKE_BATCH_SIZE_LIMIT=1
+export NRL_TQ_MOONCAKE_MAX_WORKER_THREADS=1
+export MC_TCP_ENABLE_CONNECTION_POOL=1
+
+uv run --locked --extra modelopt --extra mcore --extra nemo-gym \
+    python examples/nemo_gym/run_distillation_nemo_gym.py \
+    --config examples/modelopt/qa_distillation_nano3_megatron_tq.yaml \
+    "policy.model_name=$MODEL" \
+    "policy.tokenizer.name=$MODEL" \
+    "teacher.model_name=$MODEL" \
+    "teacher.tokenizer.name=$MODEL" \
+    "policy.quant_calib_data=$QUANT_CALIB_DATA" \
+    "data.train.data_path=$TRAIN_DATA" \
+    "distillation.max_num_steps=$MAX_NUM_STEPS" \
+    "cluster.num_nodes=$NUM_NODES" \
+    "cluster.gpus_per_node=$GPUS_PER_NODE" \
+    "checkpointing.checkpoint_dir=$OUTPUT_DIR/checkpoints" \
+    "logger.log_dir=$OUTPUT_DIR/logs" \
+    "logger.tensorboard.log_dir=$OUTPUT_DIR/tensorboard"

examples/modelopt/nano3_qaopd_tq/submit_slurm.sh

@@ -0,0 +1,67 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+usage() {
+    cat <<'EOF'
+Usage:
+  submit_slurm.sh CONTAINER ACCOUNT PARTITION SHARED_ROOT MODEL \
+      TRAIN_DATA QUANT_CALIB_DATA OUTPUT_DIR [NUM_NODES] [GPUS_PER_NODE]
+
+CONTAINER may be an Enroot .sqsh file or a registry image accepted by Pyxis.
+All local paths must be visible from every allocated node.
+EOF
+}
+
+if (( $# < 8 || $# > 10 )); then
+    usage >&2
+    exit 2
+fi
+
+CONTAINER_IMAGE=$1
+SLURM_ACCOUNT=$2
+SLURM_PARTITION=$3
+SHARED_ROOT=$4
+MODEL=$5
+TRAIN_DATA=$6
+QUANT_CALIB_DATA=$7
+OUTPUT_DIR=$8
+NUM_NODES=${9:-8}
+GPUS_PER_NODE=${10:-4}
+
+REPO_ROOT=$(git rev-parse --show-toplevel)
+mkdir -p "$OUTPUT_DIR"
+
+printf -v TRAIN_COMMAND '%q ' \
+    bash examples/modelopt/nano3_qaopd_tq/run_training.sh \
+    "$TRAIN_DATA" "$QUANT_CALIB_DATA" "$OUTPUT_DIR" "$MODEL"
+COMMAND="cd /opt/nemo-rl && $TRAIN_COMMAND"
+MOUNTS="$SHARED_ROOT:$SHARED_ROOT,$REPO_ROOT:/opt/nemo-rl"
+
+CONTAINER="$CONTAINER_IMAGE" \
+MOUNTS="$MOUNTS" \
+COMMAND="$COMMAND" \
+BASE_LOG_DIR="$OUTPUT_DIR/ray" \
+GPUS_PER_NODE="$GPUS_PER_NODE" \
+sbatch \
+    --nodes="$NUM_NODES" \
+    --account="$SLURM_ACCOUNT" \
+    --partition="$SLURM_PARTITION" \
+    --time=04:00:00 \
+    --exclusive \
+    --dependency=singleton \
+    --job-name=nano3-qaopd-tq \
+    ray.sub