Mercity/policy-classifier/05-11-progress/training - Storage Bucket

Prepare 25k pairs and train on an A100

The 25k configurations are prepared. The Modal entry point runs the exported 5k dataset through an A100 smoke test and training.

Generate and export (repo root)

The prepared generation config requests five cases per training domain and 20 conversations per case (60 cases, approximately 1,200 conversations). This leaves headroom for validation, rejected examples and deduplication. The edit stage creates whole-conversation variants with small followed-to-violated changes; its request budget is configurable. The models and pricing settings match the existing pipeline.

uv run python -X utf8 -m curator_pipeline.run curator_pipeline/config_25k.yaml

The runner performs cases → conversations → judge → edit → data export → viewer. It does not train. To rerun only the exporter against an existing completed run:

uv run python -X utf8 data.py --config curator_pipeline/config_25k.yaml

datasets/policy_25k/train.jsonl contains exactly 25,000 message–policy pairs. Validation is additional and uses separate cases, with at least one validation case per training domain. Originals and edited variants share a case split. Held-out domains, when generated in the source run, go only to test_ood.jsonl; the default config excludes them, so that file will initially be empty. This is not yet an OOD benchmark. Identical inputs are deduplicated, conflicting labels are quarantined, and stale/incomplete source artifacts are refused or excluded. The exporter errors if insufficient training data remains; it never silently fills the count with repeated samples. manifest.json records counts, labels and source hashes. rejected.jsonl records exclusions.

Labels are followed=0, violated=1, not_applicable=2. The model input is solely the message and policy. It does not include judge reasons, agent context, preceding user requests, or conversation history. The judge uses full conversation context. Requirements dependent on earlier turns may remain ambiguous to the message-only classifier.

A100 environment (Linux, repo root)

Use a separate Python 3.11/3.12 environment; the training dependency file is separate from the Curator generation dependencies. Install the CUDA-enabled PyTorch build appropriate to the host first if the environment does not already provide one.

python -m venv .venv-train
source .venv-train/bin/activate
python -m pip install -r training/requirements.txt
wandb login

Copy the repo and exported datasets/policy_25k/ directory to the A100 host. No credentials are stored in configs. W&B logging is enabled; its authentication happens on that host.

Single H100 pod: domain split, six runs (current)

All commands run from the repo root. Data: policy_v2/ (exported 27k train / 3k validation / 7,023 test_ood) re-split into policy_v2_domain/: train and validation are pooled, the 10 training domains keep about 10% of each domain's cases (whole cases, so edits stay with their originals) as validation, and the 2 held-out domains (insurance_claims, travel_booking) are the test set, copied unchanged.

# 1. Environments. System Python 3.12 already has torch 2.8.0+cu128.
python -m pip install -r training/requirements.txt hf_transfer   # transformers 4.57.3: ModernBERT, Qwen3
wandb login
# Qwen3.5 needs transformers 5; a separate venv reuses the system torch (no torch reinstall).
uv venv --system-site-packages --python /usr/local/bin/python .venv-qwen35
VIRTUAL_ENV=.venv-qwen35 uv pip install --no-deps "transformers==5.20.2" "peft==0.21.2"
VIRTUAL_ENV=.venv-qwen35 uv pip install --exclude-newer 2026-10-22 "huggingface-hub>=1.31,<3" "tokenizers>=0.23.1,<0.24" typer "regex>=2025.10.22"

# 2. Models (cached in $HF_HOME).
hf download answerdotai/ModernBERT-large && hf download Qwen/Qwen3-0.6B && hf download Qwen/Qwen3.5-0.8B

# 3. Domain split (refuses to overwrite an existing output folder).
python -I scripts/resplit_by_domain.py policy_v2 policy_v2_domain

# 4. Preflight (no GPU): every config should print "Preflight passed".
python -m training.train_modernbert --config training/modernbert_large_30k.yaml --preflight
python -m training.train_qwen --config training/qwen3_06b_30k.yaml --preflight
.venv-qwen35/bin/python -m training.train_qwen35 --config training/qwen35_08b_30k.yaml --preflight

# 5. One run = smoke test (picks batch settings, writes resolved.yaml) -> training -> W&B upload:
CUDA_VISIBLE_DEVICES=0 python -m training.smoke --config training/modernbert_large_30k.yaml
CUDA_VISIBLE_DEVICES=0 python -m training.train_modernbert --config checkpoints/modernbert-large-30k/resolved.yaml > modernbert.log 2>&1
python scripts/log_checkpoints_wandb.py checkpoints/modernbert-large-30k   # checkpoints + reports as W&B artifacts

# Or the queues used here (one GPU, runs strictly one after another; each waits for the pid given):
nohup scripts/run_queue.sh > queue.log 2>&1 &                     # ModernBERT then Qwen3, 1.5 epochs
nohup scripts/run_queue_2ep.sh <pid> > queue_2ep.log 2>&1 &       # both again, 2 epochs
nohup scripts/run_queue_qwen35.sh <pid> > queue_qwen35.log 2>&1 & # Qwen3.5-0.8B, 1.5 then 2 epochs

# 6. Numbers for every finished run (per-checkpoint validation F1, final validation/test_ood, natural vs edited):
python -I scripts/collect_results.py checkpoints/*/

Config	Model	Epochs	max_length
training/modernbert_large_30k.yaml / _2ep.yaml	ModernBERT-large	1.5 / 2	2,400
training/qwen3_06b_30k.yaml / _2ep.yaml	Qwen3-0.6B	1.5 / 2	2,352
training/qwen35_08b_30k.yaml / _2ep.yaml	Qwen3.5-0.8B, text decoder only	1.5 / 2	2,360

All six share LoRA rank 32, alpha 64, dropout 0.05, effective batch 64, learning rate 2e-4. The smoke test may replace batch 64 x 1 with e.g. 16 x 4 (same effective batch) when that is faster and fits; the training run then reads resolved.yaml. With gradient accumulation the Trainer's logged training loss is summed over micro-batches (divide by the accumulation steps); evaluation loss is unaffected. final/ is the last checkpoint, not the best one; compare the four quarter checkpoints in W&B. Every run, with its paths and W&B link, is listed in CATALOGUE.md.

Train both classifiers on a 2x A100 pod

Two scripts share all training machinery in training/engine.py; each only loads its backbone, picks its LoRA target modules and classification head, and builds its collator.

Script	Config	Model / head	LoRA targets	max_length
training/train_modernbert.py	training/modernbert_large_30k.yaml	ModernBERT-large, mean-pooled head + classifier	Wqkv, Wi, Wo	2,400
training/train_qwen.py	training/qwen3_06b_30k.yaml	Qwen3-0.6B, score on the last non-pad token	all attention + MLP projections	2,352

max_length is the longest tokenized input in datasets/policy_30k (1,397 ModernBERT / 1,347 Qwen tokens) plus 1,000, rounded up to a multiple of 8. Shared hyperparameters: LoRA rank 32, alpha 64, dropout 0.05 (use_dora: false; set true for DoRA), batch 64 per GPU with no accumulation, 1.5 epochs, learning rate 2e-4 with linear decay and 5% warmup, weight decay 0.01, gradient clipping 1.0, BF16, gradient checkpointing on. Copy datasets/policy_30k/ to the pod (datasets are not in Git).

# Check config + data without loading a model or touching the GPU:
python -m training.train_modernbert --config training/modernbert_large_30k.yaml --preflight
python -m training.train_qwen --config training/qwen3_06b_30k.yaml --preflight
# Optional: 10-example smoke test + worst-length batch profiling (writes resolved.yaml):
CUDA_VISIBLE_DEVICES=0 python -m training.smoke --config training/modernbert_large_30k.yaml
# One model per GPU, in parallel (each keeps an effective batch of 64):
CUDA_VISIBLE_DEVICES=0 nohup python -m training.train_modernbert --config training/modernbert_large_30k.yaml > modernbert.log 2>&1 &
CUDA_VISIBLE_DEVICES=1 nohup python -m training.train_qwen --config training/qwen3_06b_30k.yaml > qwen.log 2>&1 &

Do not launch one script with torchrun across both GPUs unless you halve batch_size; the effective batch would otherwise double to 128.

Checkpoints: adapter weights plus the trained classification head only (no optimizer or full model), at 1/4, 2/4, 3/4 and the end of training, in checkpoints/<run>/checkpoint-<k>of4-step<N>/, with final/ equal to the last one. Validation runs at each of those points and logs to W&B (policy-classifier). After training, validation_predictions.jsonl holds every validation row with its predicted label and class probabilities; validation_report.json has the metrics, confusion matrix and scikit-learn classification report; W&B gets the confusion matrix and up to 500 errors.

Metrics come only from scikit-learn (precision_recall_fscore_support, accuracy_score, macro f1_score, confusion_matrix) with all three labels fixed, so a missing class never shifts the columns. A unit test checks them against hand-computed precision, recall, F1, accuracy and macro F1.

Score any checkpoint on any split later (GPU if present):

python -m training.evaluate --checkpoint checkpoints/modernbert-large-30k/checkpoint-2of4-step<N> --data datasets/policy_30k/validation.jsonl
python -m training.predict --checkpoint checkpoints/qwen3-0.6b-30k/final --message "..." --policy "..."

evaluate also reads older DoRA checkpoints, such as the 5k ModernBERT final/ on the Modal volume.

Qwen3.5-0.8B is not supported by the pinned transformers 4.57.3 (it is a hybrid-attention model), so train_qwen.py refuses it; train_qwen35.py trains it from the separate transformers-5 environment (.venv-qwen35, see above). Qwen3.5 checkpoints are vision-language models: only the text decoder is loaded, through Hugging Face's Qwen3_5TextForSequenceClassification (last non-pad token pooling, as in Qwen3). Three of every four decoder layers are Gated DeltaNet linear attention; LoRA covers their in_proj_qkv, in_proj_z, out_proj, the full-attention q/k/v/o_proj and the MLP. Without the optional flash-linear-attention kernels, transformers uses its slower PyTorch implementation of those layers (same results). engine.py works with both transformers versions (warmup_ratio became warmup_steps in 5.x; the warmup step count is identical).

Older configurations

training/modernbert_large.yaml, modernbert_base.yaml (25k plan) and the 5k Modal configs were migrated to the same keys (lora_rank, use_dora: true) and run with the matching script, e.g. python -m training.train_modernbert --config training/modernbert_large.yaml.

ModernBERT is an encoder, so there are no chat system/user roles. A fixed task instruction prefixes the message, and the pretrained tokenizer inserts its existing separator between that field and the policy. No new token or frozen random embedding is introduced. Training and inference share the same format. Truncation preserves the policy and truncates only the message field.

The model uses Hugging Face's standard ModernBERT classification implementation and its cross-entropy loss. DoRA targets the attention/MLP projections Wqkv, Wi, and Wo; the head and classifier are fully trained and saved alongside the adapter. The backbone is loaded in BF16, but every trainable weight (LoRA and the head) is kept in float32 (upcast_trainable_parameters in training/engine.py); in BF16 most of Adam's small head updates rounded to zero. The seed is set before the model is built, so the head and LoRA start from the same values on every run with the same seed. Validation and test reports (*_report.json) include by_source: the same metrics for natural pairs and edited pairs separately, also logged to W&B as <split>_natural/... and <split>_edited/.... Training refuses to start when WANDB_MODE is offline unless the config sets allow_offline_wandb: true. Engine tests (tiny models on CPU) need torch: D:/venvs/train-cpu/Scripts/python.exe -m unittest tests.test_training_engine. The best validation macro-F1 checkpoint is retained. Metrics include accuracy and per-class precision, recall and F1 so that majority-class accuracy cannot hide poor violation detection. Metric calculations use scikit-learn's precision_recall_fscore_support, f1_score and accuracy_score, with all three label ids explicit and zero_division=0. W&B receives metrics and a sample of evaluation predictions; resolved configs and metrics are also saved locally.

The final/ artifact contains the adapter, classification head, tokenizer and input-format recipe. The recipe records the resolved base-model revision so inference reloads the same pretrained weights. Resume an interrupted run with --resume-from-checkpoint checkpoints/.../checkpoint-N. Track executed experiments and W&B run links in CATALOGUE.md.

Inference

python -m training.predict --checkpoint checkpoints/modernbert_large_dora32/final \
  --message "I cannot provide a diagnosis; please discuss your symptoms with your doctor." \
  --policy "Never provide a medical diagnosis."

Implementation references

• ModernBERT-large model card
• ModernBERT sequence-classification implementation, Transformers 4.57.3
• PEFT DoRA configuration

Offline tests: python -X utf8 -m unittest discover -s tests -v. Full model loading, adapter serialization and the A100 training step still need a smoke test in the GPU environment.

Modal A100: 5k smoke test and training

Install the Modal client locally and use the authenticated rishik00 profile. That workspace must contain the wandb Secret with WANDB_API_KEY. The image uploads the training package and exported 5k data only; it does not upload .env, generation caches, or the rest of the repository.

uv pip install --python D:/venvs/mercity-train_policy_classifier/Scripts/python.exe "modal>=1,<2"
D:/venvs/mercity-train_policy_classifier/Scripts/python.exe -X utf8 -m modal run --detach --profile rishik00 -m training.modal_app --epochs 1.5
# Smoke test only:
D:/venvs/mercity-train_policy_classifier/Scripts/python.exe -X utf8 -m modal run --detach --profile rishik00 -m training.modal_app --smoke-only

The remote job uses one A100 80GB, BF16, DoRA rank 32, and the original three-class classifier. It first trains on ten distinct examples, checks finite gradients and adapter save/reload, then benchmarks microbatches 4, 8, 16, 32, and 64 with checkpointing on and off. Larger profiling batches repeat the longest real input solely to measure memory. The ten examples are not duplicated into the actual dataset. Candidate settings must leave 15% of GPU memory free, and the fastest safe candidate is selected while preserving effective batch 64. The context cap remains 8192; dynamic padding uses actual input lengths. Full training starts from fresh pretrained weights after the smoke test succeeds and uses exactly the exported 5000 training pairs.

Artifacts persist on Modal Volume policy-nli-training in the directory recorded by scratch/latest_modal_training.txt: status.json, smoke_summary.json, resolved.yaml, hardware_metrics.json, Hugging Face metrics/checkpoints, and final/. Smoke and training each log to W&B. The full run is skipped if the smoke test fails. Download outputs with:

D:/venvs/mercity-train_policy_classifier/Scripts/python.exe -X utf8 -m modal volume get --profile rishik00 policy-nli-training <run-id> checkpoints/<run-id>

Qwen3-0.6B comparison on the same 5k pairs

D:/venvs/mercity-train_policy_classifier/Scripts/python.exe -X utf8 -m training.train_qwen --config training/qwen3_06b_5k.yaml --preflight
D:/venvs/mercity-train_policy_classifier/Scripts/python.exe -X utf8 -m modal run --detach --profile rishik00 -m training.modal_app --config training/qwen3_06b_5k.yaml --epochs 1.5

This uses Qwen/Qwen3-0.6B with Hugging Face's standard three-class sequence-classification head and ordinary cross-entropy. It preserves the exact 5k train/validation split, seed, learning rate, epoch count, and DoRA rank from the ModernBERT comparison. The smoke test chooses safe microbatch settings independently for this architecture. DoRA targets q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, and down_proj; the score classification head is fully trained and saved. The standard Qwen classifier pools the last non-padding token. Generation KV caching is disabled, and the model uses its existing padding token (or existing EOS token if needed).

The text has the same instruction, assistant-message heading, and policy heading as ModernBERT. Qwen has no pretrained pair separator, so ordinary blank-line text separates the fields. No special token or embedding is added, and no chat template, reasoning generation, or language-model label generation is used. Shared tokenization preserves the policy when truncating long inputs; the exported input-format recipe and tokenizer are used by training.predict.

Architecture references:

• https://huggingface.co/Qwen/Qwen3-0.6B
• https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/qwen3/modeling_qwen3.py
• https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/modeling_layers.py