env_artifacts/README.md

Lightning AI H100 environment — reproduction artifacts

Files in this folder capture the exact, working stack we landed on for corp_gym SFT/RLVR training. Use them to bring up a fresh Lightning AI Studio (or any Linux H100 with driver >= 570, CUDA 12.8 runtime, Python 3.12).

Contents

• setup_lightning_h100.sh — curated install script (recommended). Pins torch 2.7.1 + cu128, xformers 0.0.31.post1, FlashAttention 2.8.0.post2, unsloth 2026.4.8, trl 0.24.0, peft 0.19.1, torchao 0.17.0, transformers 5.5.0, bitsandbytes 0.49.2, datasets 4.3.0, accelerate 1.13.0, plus matplotlib>=3.10 and numpy<3 for the plotting step.
• requirements_lightning_h100.txt — the same pins in a single pip file (grouped; read the top comments for the ordered install commands).
• requirements_frozen.txt — full pip freeze of the exact environment (treat as a ground truth, not a reinstall recipe — it does not know about --index-url or wheel URLs).

Recommended Lightning Studio environment

• GPU: 1x NVIDIA H100 80GB
• Runtime: Linux with CUDA 12.8 compatible driver (>= 570.x)
• Python: 3.12 (base Lightning image is fine)
• Start from a fresh Studio and run bash env_artifacts/setup_lightning_h100.sh

This is the fastest conflict-free path for Unsloth in this repo because it pins a known-good torch/xformers/flash-attn intersection.

Tested against

• GPU: NVIDIA H100 80GB HBM3 (sm_90), driver 570.148.08
• CUDA runtime: 12.8
• Python: 3.12.11

Why these pins (not the defaults from pyproject.toml)

The factory Lightning image ships torch 2.5.0 (cu121). That trips a chain of compat breaks:

1. torchao >= 0.13 references torch.int1 (added in torch 2.6).
2. torchao >= 0.17 references torch.utils._pytree.register_constant (added in torch 2.7).
3. peft 0.19.1 requires torchao > 0.16 if any torchao is importable.
4. unsloth_zoo 2026.4.9 hard-requires torchao >= 0.13.
5. flash-attn prebuilt wheels only cover torch 2.4–2.8; 2.9/2.10/2.11 must be compiled from source against nvcc (fragile and slow).
6. xformers couples tightly to both torch and a narrow flash-attn range.

Torch 2.7.1 + xformers 0.0.31.post1 + flash-attn 2.8.0.post2 is the intersection where all four constraints (int1, register_constant, peft>torchao>0.16, prebuilt FA2 wheel) are satisfied.

One-shot bring-up

# In a fresh Lightning Studio with this repo checked out:
cd corp_gym
bash env_artifacts/setup_lightning_h100.sh
pip install -e .          # picks up project deps without touching the pinned torch stack

# Per-session exports (stick these in ~/.bashrc if you want):
export HF_HUB_ENABLE_HF_TRANSFER=1
export TOKENIZERS_PARALLELISM=false
export TRANSFORMERS_VERBOSITY=warning

huggingface-cli login   # once per container

Resuming training from the previously-trained adapters

The prior run pushed both adapters to Navigam/*:

• Navigam/corp-env-sft-qwen2.5-7b

You can pull them locally on a new box with:

huggingface-cli download Navigam/corp-env-sft-qwen2.5-7b --local-dir outputs/sft_adapter

Then for a direct 14B RLVR run:

python training/train_rlvr.py \
  --model Qwen/Qwen3-14B-Instruct \
  --adapter outputs/sft_qwen3_14b \
  --examples data/processed/e1_m1_clean.jsonl,data/processed/h1_seed_clean.jsonl \
  --output outputs/rlvr_qwen3_14b \
  --rounds 3 --n-samples 8 --max-prompts 128 \
  --stats-file results/runs/rlvr_stats.jsonl \
  --push-to-hub <your-user-or-org>/corp-env-rlvr-qwen3-14b

Notes / gotchas

• First time only on any new box: unsloth writes patched-trainer classes into corp_gym/unsloth_compiled_cache/. Delete that folder if you ever change TRL/unsloth versions to avoid stale compiled patches.
• Flash-attn fp32 workaround: this is now legacy-only for training/train_grpo.py. It is disabled by default and can be enabled only when needed via CORP_ENABLE_FA2_BF16_PATCH=1.
• training/train_rlvr.py now runs without this monkey patch on modern stacks.
• max_prompt_length filter: both legacy GRPO and RLVR scripts tokenise every prompt up-front and drop rows whose chat-template-encoded length exceeds 0.9 * max_prompt_length (long H1 trajectories otherwise produce a causal mask / attention mask size mismatch at generation time).