arianraje/mimo-7b-gdn-opd-predecay-1337m-step8676

Bootstrap the live MiMo OPD control on a new machine

This is the portable recovery path for the unchanged H=16K MiMo control. The private full-state checkpoint is:

arianraje/mimo-7b-gdn-opd-resume-step2350

It resumes at step 2350 / 353,104,411 generated tokens, with two trainer ranks, one sampler, horizon 16,384, and W&B ID mimo-rugfix-h200x3-ext200-decay0-v1-20260806. The target remains 403,158,000 tokens. Do not launch it while another job owns the same output directory.

1. Clone and create the environment

git clone git@github.com:arianraje/Linearization.git
cd Linearization
python3.10 -m venv /fast/path/mimo_env
source /fast/path/mimo_env/bin/activate
python -m pip install --upgrade pip wheel setuptools
python -m pip install -r requirements-mimo-resume.txt

Hopper builds may require TORCH_CUDA_ARCH_LIST=9.0 and --no-build-isolation for causal-conv1d or flash-linear-attention. tilelang==0.1.12 is mandatory: FLA GDN backward fails on Hopper without it. The exact validated software versions are pinned in requirements-mimo-resume.txt.

2. Download the private recovery state and public teacher

Authenticate with an arianraje token, or a read token explicitly granted access to the private recovery repository:

hf auth login
export MIMO_STORAGE=/fast/path/mimo
mkdir -p "$MIMO_STORAGE/checkpoints" "$MIMO_STORAGE/models" "$MIMO_STORAGE/data"
hf download arianraje/mimo-7b-gdn-opd-resume-step2350 \
  --local-dir "$MIMO_STORAGE/checkpoints/mimo-control-step2350"
hf download XiaomiMiMo/MiMo-7B-RL-0530 \
  --local-dir "$MIMO_STORAGE/models/MiMo-7B-RL-0530"

The recovery repository is roughly 109 GB. Keep at least 230 GB free because atomic optimizer saves temporarily need another approximately 100 GB.

Validate the download without loading the 100 GB optimizer payload into RAM:

python scripts/validate_mimo_resume_checkpoint.py \
  "$MIMO_STORAGE/checkpoints/mimo-control-step2350" \
  --expected-step 2350 --expected-tokens 353104411 --expected-world 2 \
  --manifest configs/mimo_resume_step2350.json

The BF16 model files in this rolling layout are the branch seed, while the optimizer shards contain the current FP32 master parameters. This is intentional: ShardedMasterAdamW.load_state_dict copies the step-2350 masters into every BF16 model parameter before the first resumed rollout.

3. Recreate the two deterministic data directories

The prompt set is small and the 32K validation mix is deterministic, so they are rebuilt rather than uploaded. Both builders need internet. Do not run the 32K mix builder inside a memory-limited login shell; give it at least 64 GB.

export HF_HUB_OFFLINE=0
python -m src.data.build_stage3_prompts \
  --output "$MIMO_STORAGE/data/stage3_prompts_v1"
python -m src.data.build_stage2b_mix \
  --output "$MIMO_STORAGE/data/stage2b_mix_mimo_v1" \
  --tokenizer "$MIMO_STORAGE/models/MiMo-7B-RL-0530"

Acceptance checks:

• stage3_prompts_v1: 60,000 reasoning prompts, 15,011 general prompts, 129 contaminated prompts removed.
• stage2b_mix_mimo_v1: separator token 151643; approximately 294.09M train and 6.00M validation tokens at 32,768; seed 1234.

4. Materialize a machine-local resume config

python scripts/make_mimo_resume_config.py \
  --checkpoint "$MIMO_STORAGE/checkpoints/mimo-control-step2350" \
  --teacher "$MIMO_STORAGE/models/MiMo-7B-RL-0530" \
  --prompts "$MIMO_STORAGE/data/stage3_prompts_v1" \
  --validation-data "$MIMO_STORAGE/data/stage2b_mix_mimo_v1" \
  --output "$MIMO_STORAGE/mimo_resume.yaml"

Inspect the resulting paths and confirm resume: true, train_gen_tokens: 403158000, horizon_max: 16384, and the unchanged W&B ID. Set W&B credentials outside the repository.

5. Launch inside three H200-class GPUs

The observed control high-water mark was about 125.6 GB per trainer, so this recipe requires H200-class memory. The topology is two trainer ranks plus one vLLM sampler. First obtain a three-GPU allocation, then run:

export MIMO_REPO_ROOT="$PWD"
export MIMO_CHECKPOINT="$MIMO_STORAGE/checkpoints/mimo-control-step2350"
export MIMO_CONFIG="$MIMO_STORAGE/mimo_resume.yaml"
export HF_HOME="$MIMO_STORAGE/hf_cache"
export HF_HUB_OFFLINE=1
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export OMP_NUM_THREADS=4
bash scripts/run_mimo_resume_in_allocation.sh

On Slurm, request three H200s, at least 16 CPUs, and enough wall time for the remaining approximately 50.05M tokens (about 7.4 training hours on the validated Empire topology, plus startup/checkpoint margin). The launcher refuses sub-130-GB GPUs and uses an output writer lock.

Scientific state to preserve

• The existing trajectory is the unchanged control.
• Latest NIAH multikey-32K at 340M: 0.800 (n=100, watch, not recovery).
• GSM8K strict/flexible at 340M: 0.55876 / 0.66187; strict remains below the 250M best.
• Preserve the armed 380M probe and final 400M evaluation.
• The leading NIAH mechanism finding is prefix poisoning: the teacher chose the correct clean-prefix clause 8/8 times, but only 3/8 after corrupted student reasoning. Do not silently alter the control recipe.

Read docs/HANDOFF_MIMO_OPD_LIVE_NIAH.md completely before taking over live operations. Scheduler job IDs and systemd/tmux watcher state in that handoff belong to the original machine and must not be replayed on the new one.

Source Git commit: 1b82f7c6327e4f17d99ff7b1dc76a57789bf67a4