Mercity/policy-classifier/03-11-progress - 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.

Setup on this pod (1x H100, Python 3.12)

The repo is laid out as a training/ package with data under datasets/. Commands run from the repo root.

uv venv --python /usr/bin/python3 .venv-train
source .venv-train/bin/activate
uv pip install "torch==2.8.0" --index-url https://download.pytorch.org/whl/cu128
uv pip install -r training/requirements.txt hf_transfer "huggingface_hub[cli]" "torch==2.8.0"
hf download answerdotai/ModernBERT-large --exclude "*.onnx" "onnx/*"
hf download Qwen/Qwen3-0.6B
wandb login   # or export WANDB_MODE=offline and sync later with: wandb sync wandb/offline-run-*

Domain split: 10 training domains, 2 held-out test domains

Builds datasets/policy_30k_domains from datasets/policy_30k. Held-out domains go to test_ood.jsonl; validation is one whole case per training domain. Training then scores both validation and test.

python -m training.split_by_domain --source datasets/policy_30k --out datasets/policy_30k_domains \
  --test-domains insurance_claims travel_booking --val-cases-per-domain 1 --seed 7
python -m training.train_modernbert --config training/modernbert_large_30k_domains.yaml --preflight
python -m training.train_qwen --config training/qwen3_06b_30k_domains.yaml --preflight
# Both fit on one H100 together (~41 GB):
CUDA_VISIBLE_DEVICES=0 nohup python -m training.train_modernbert --config training/modernbert_large_30k_domains.yaml > modernbert_domains.log 2>&1 &
CUDA_VISIBLE_DEVICES=0 nohup python -m training.train_qwen --config training/qwen3_06b_30k_domains.yaml > qwen_domains.log 2>&1 &

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); train_qwen.py refuses it. Using it requires upgrading transformers and re-checking the head.

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 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