Fit the TrainerRank layout cost model (coefficient version 2) with a calibration campaign

dev/trainer_rank_cost_calibration_2gpu.sky.yaml

@@ -0,0 +1,97 @@
+# Cost-model calibration: 2-GPU cells (2x H200, Kubernetes).
+#
+# TP2 x CP1 GRPO/heterogeneous cells (the first TP calibration evidence) and
+# TP1 x CP2 Ellavox groups (real data; the corpus is git-ignored and mounted).
+# Even Ellavox groups are calibration rows, odd groups are held out for
+# validation by the fit (see dev/trainer_rank_cost_fit.py).
+#
+# Usage:
+#   sky launch -c tr-cost-2gpu dev/trainer_rank_cost_calibration_2gpu.sky.yaml --env RUN_ID=<id>
+
+name: trainer-rank-cost-calibration-2gpu
+
+workdir: .
+
+file_mounts:
+  /tmp/trainer_rank_ellavox_corpus.json: dev/_trainer_rank_ellavox_qwen35_4b_tokens.json
+
+envs:
+  RUN_ID: unset
+  REPEAT: "8"
+  ELLAVOX_GROUPS: "0 1 2 3 4 5 6 7"
+  # "main" = TP2 GRPO/hetero + Ellavox CP2; "cp2-hetero" = GRPO/hetero at CP2
+  # and heterogeneous variants at TP2 (second campaign).
+  CELL_SET: main
+
+resources:
+  cloud: kubernetes
+  accelerators: H200:2
+  cpus: 24+
+  memory: 192+
+  image_id: docker:docker.io/bradhiltonnw/art-gpu@sha256:f93fda4032bf37f16d9c639c51daf73e1a884df97d14baabf349868831321bdc
+
+setup: |
+  set -euo pipefail
+  export CC=gcc CXX=g++
+  uv sync --frozen --extra megatron --group dev
+  uv sync --project megatron_runtime --extra cuda12 --group test --frozen
+  uv pip install --python megatron_runtime/.venv/bin/python --no-deps --editable .
+
+run: |
+  set -euo pipefail
+  : "${RUN_ID:?RUN_ID is required}"
+  root="scratch/trainer_rank_cost_calibration/${RUN_ID}"
+  mkdir -p "${root}"
+  cp /tmp/trainer_rank_ellavox_corpus.json dev/_trainer_rank_ellavox_qwen35_4b_tokens.json
+  export ART_MEGATRON_DATA_PARALLEL_SIZE=1
+  export ART_MEGATRON_PIPELINE_MODEL_PARALLEL_SIZE=1
+  export TOKENIZERS_PARALLELISM=false
+  export PYTHONUNBUFFERED=1
+  uv run --no-sync python dev/trainer_rank_landing_acceptance.py --phase contract
+
+  cell() {  # tp cp cell model layers [group]
+    ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE="$1" ART_MEGATRON_CONTEXT_PARALLEL_SIZE="$2" \
+    timeout 7200s megatron_runtime/.venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=2 \
+      dev/trainer_rank_landing_acceptance.py --phase cost-calibrate \
+      --cell "$3" --model "$4" --layers "$5" --group "${6:-0}" --repeat "${REPEAT}" \
+      --evidence "${root}/evidence.jsonl" 2>&1 | tee "${root}/tp$1-cp$2-$3-$5-g${6:-0}.log" | { grep -E "cost-calibrate|LANDING|Traceback" || true; }
+  }
+  # pipefail carries torchrun's status through tee and the grep filter; a
+  # failed cell is recorded and the campaign exits nonzero at the end.
+  failures=0
+  run_cell() { cell "$@" || { failures=$((failures + 1)); echo "CELL FAILED: $*"; }; }
+  if [ "${CELL_SET}" = "main" ]; then
+    run_cell 2 1 cal-grpo-g8   Qwen/Qwen3.5-4B 0
+    run_cell 2 1 cal-grpo-g16  Qwen/Qwen3.5-4B 0
+    run_cell 2 1 cal-grpo-g4x4 Qwen/Qwen3.5-4B 0
+    run_cell 2 1 cal-hetero    Qwen/Qwen3.5-4B 0
+    run_cell 2 1 cal-grpo-g8   Qwen/Qwen3-4B   0
+    for g in ${ELLAVOX_GROUPS}; do
+      run_cell 1 2 cal-ellavox Qwen/Qwen3.5-4B 0 "$g"
+    done
+  else
+    run_cell 1 2 cal-grpo-g8   Qwen/Qwen3.5-4B 0
+    run_cell 1 2 cal-grpo-g16  Qwen/Qwen3.5-4B 0
+    run_cell 1 2 cal-grpo-g4x4 Qwen/Qwen3.5-4B 0
+    run_cell 1 2 cal-hetero    Qwen/Qwen3.5-4B 0
+    run_cell 1 2 cal-hetero2   Qwen/Qwen3.5-4B 0
+    run_cell 1 2 cal-hetero3   Qwen/Qwen3.5-4B 0
+    run_cell 1 2 cal-grpo-g8   Qwen/Qwen3-4B   0
+    run_cell 2 1 cal-hetero2   Qwen/Qwen3.5-4B 0
+    run_cell 2 1 cal-hetero3   Qwen/Qwen3.5-4B 0
+  fi
+  echo "cost-calibration 2gpu (${CELL_SET}): done (${failures} failed cells)"
+  [ "${failures}" -eq 0 ]
+
+config:
+  kubernetes:
+    pod_config:
+      spec:
+        schedulerName: binpack-scheduler
+        activeDeadlineSeconds: 42000
+        containers:
+          - name: ray-node
+            imagePullPolicy: Always
+            env:
+              - name: UV_LINK_MODE
+                value: copy

dev/trainer_rank_cost_calibration_cp4.sky.yaml

@@ -0,0 +1,100 @@
+# Cost-model calibration: CP4 cells (4x H200, Kubernetes).
+#
+# Times every mandatory candidate layout of each cell through the public API
+# (forward + backward, active LoRA) with the test-only anchor hook, recording
+# max-rank compile-free timings and layout features for the fit
+# (dev/trainer_rank_cost_fit.py). Cells: the sealed GRPO win shape on the
+# 2-layer GDN model (the one cell with retained research timings, for
+# cross-checking), the same shape at full height where memory allows, a g16
+# and a 4x4 GRPO shape, the heterogeneous control, and the attention model
+# (Qwen3-4B) at 2 layers for the non-GDN path.
+#
+# Usage:
+#   sky launch -c tr-cost-cp4 dev/trainer_rank_cost_calibration_cp4.sky.yaml --env RUN_ID=<id>
+
+name: trainer-rank-cost-calibration-cp4
+
+workdir: .
+
+# The corpus is git-ignored (customer-derived tokens); mounted for the Ellavox
+# CP4 cells and never committed.
+file_mounts:
+  /tmp/trainer_rank_ellavox_corpus.json: dev/_trainer_rank_ellavox_qwen35_4b_tokens.json
+
+envs:
+  RUN_ID: unset
+  REPEAT: "8"
+  # "main" = GRPO/hetero/attention cells; "hetero-ellavox" = heterogeneous
+  # variants and Ellavox groups (second campaign).
+  CELL_SET: main
+  ELLAVOX_GROUPS: "0 1 2 3 4 5 6 7"
+
+resources:
+  cloud: kubernetes
+  accelerators: H200:4
+  cpus: 32+
+  memory: 256+
+  image_id: docker:docker.io/bradhiltonnw/art-gpu@sha256:f93fda4032bf37f16d9c639c51daf73e1a884df97d14baabf349868831321bdc
+
+setup: |
+  set -euo pipefail
+  export CC=gcc CXX=g++
+  uv sync --frozen --extra megatron --group dev
+  uv sync --project megatron_runtime --extra cuda12 --group test --frozen
+  uv pip install --python megatron_runtime/.venv/bin/python --no-deps --editable .
+
+run: |
+  set -euo pipefail
+  : "${RUN_ID:?RUN_ID is required}"
+  root="scratch/trainer_rank_cost_calibration/${RUN_ID}"
+  mkdir -p "${root}"
+  export ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE=1
+  export ART_MEGATRON_CONTEXT_PARALLEL_SIZE=4
+  export ART_MEGATRON_DATA_PARALLEL_SIZE=1
+  export ART_MEGATRON_PIPELINE_MODEL_PARALLEL_SIZE=1
+  export TOKENIZERS_PARALLELISM=false
+  export PYTHONUNBUFFERED=1
+  uv run --no-sync python dev/trainer_rank_landing_acceptance.py --phase contract
+
+  cp /tmp/trainer_rank_ellavox_corpus.json dev/_trainer_rank_ellavox_qwen35_4b_tokens.json
+  cell() {  # cell model layers [group]
+    timeout 7200s megatron_runtime/.venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=4 \
+      dev/trainer_rank_landing_acceptance.py --phase cost-calibrate \
+      --cell "$1" --model "$2" --layers "$3" --group "${4:-0}" --repeat "${REPEAT}" \
+      --evidence "${root}/evidence.jsonl" 2>&1 | tee "${root}/$1-$3-g${4:-0}.log" | { grep -E "cost-calibrate|LANDING|Traceback" || true; }
+  }
+  # pipefail carries torchrun's status through tee and the grep filter; a
+  # failed cell is recorded and the campaign exits nonzero at the end.
+  failures=0
+  run_cell() { cell "$@" || { failures=$((failures + 1)); echo "CELL FAILED: $*"; }; }
+  if [ "${CELL_SET}" = "main" ]; then
+    run_cell cal-grpo-g8-long Qwen/Qwen3.5-4B 2
+    run_cell cal-grpo-g8      Qwen/Qwen3.5-4B 0
+    run_cell cal-grpo-g16     Qwen/Qwen3.5-4B 0
+    run_cell cal-grpo-g4x4    Qwen/Qwen3.5-4B 0
+    run_cell cal-hetero       Qwen/Qwen3.5-4B 0
+    run_cell cal-grpo-g8-long Qwen/Qwen3-4B   2
+    run_cell cal-grpo-g8      Qwen/Qwen3-4B   0
+  else
+    run_cell cal-hetero2 Qwen/Qwen3.5-4B 0
+    run_cell cal-hetero3 Qwen/Qwen3.5-4B 0
+    run_cell cal-hetero2 Qwen/Qwen3-4B   0
+    for g in ${ELLAVOX_GROUPS}; do
+      run_cell cal-ellavox Qwen/Qwen3.5-4B 0 "$g"
+    done
+  fi
+  echo "cost-calibration cp4 (${CELL_SET}): done (${failures} failed cells)"
+  [ "${failures}" -eq 0 ]
+
+config:
+  kubernetes:
+    pod_config:
+      spec:
+        schedulerName: binpack-scheduler
+        activeDeadlineSeconds: 42000
+        containers:
+          - name: ray-node
+            imagePullPolicy: Always
+            env:
+              - name: UV_LINK_MODE
+                value: copy

dev/trainer_rank_cost_calibration_local.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+# Cost-model calibration: single-GPU cells (TP1 x CP1) on a local H200.
+#
+# Usage: dev/trainer_rank_cost_calibration_local.sh <evidence.jsonl> [repeat]
+# Requires the managed Megatron runtime venv (megatron_runtime/.venv) and, for
+# the Ellavox guardrail groups, the git-ignored corpus in dev/.
+set -euo pipefail
+evidence=${1:?evidence jsonl path}
+repeat=${2:-8}
+export PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false
+export ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE=1 ART_MEGATRON_CONTEXT_PARALLEL_SIZE=1
+export ART_MEGATRON_DATA_PARALLEL_SIZE=1 ART_MEGATRON_PIPELINE_MODEL_PARALLEL_SIZE=1
+logdir="$(dirname "${evidence}")"
+
+cell() {  # cell model layers [group]
+  timeout 7200s megatron_runtime/.venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=1 \
+    dev/trainer_rank_landing_acceptance.py --phase cost-calibrate \
+    --cell "$1" --model "$2" --layers "$3" --group "${4:-0}" --repeat "${repeat}" \
+    --evidence "${evidence}" 2>&1 | tee "${logdir}/tp1-cp1-$1-$3-g${4:-0}.log" \
+    | { grep -E "cost-calibrate|LANDING|Traceback" || true; }
+}
+# pipefail carries torchrun's status through tee and the grep filter; a failed
+# cell is recorded and the script exits nonzero at the end.
+failures=0
+run_cell() { cell "$@" || { failures=$((failures + 1)); echo "CELL FAILED: $*"; }; }
+run_cell cal-grpo-g8-long Qwen/Qwen3.5-4B 2
+run_cell cal-grpo-g8      Qwen/Qwen3.5-4B 0
+run_cell cal-grpo-g16     Qwen/Qwen3.5-4B 0
+run_cell cal-grpo-g4x4    Qwen/Qwen3.5-4B 0
+run_cell cal-hetero       Qwen/Qwen3.5-4B 0
+run_cell cal-hetero2      Qwen/Qwen3.5-4B 0
+run_cell cal-hetero3      Qwen/Qwen3.5-4B 0
+run_cell cal-grpo-g8-long Qwen/Qwen3-4B   2
+run_cell cal-grpo-g8      Qwen/Qwen3-4B   0
+run_cell cal-hetero2      Qwen/Qwen3-4B   0
+if [ -f dev/_trainer_rank_ellavox_qwen35_4b_tokens.json ]; then
+  for g in 0 1 2 3 4 5 6 7; do
+    run_cell cal-ellavox Qwen/Qwen3.5-4B 0 "$g"
+  done
+fi
+echo "cost-calibration local: done (${failures} failed cells)"
+[ "${failures}" -eq 0 ]