MODULUS Hyperball Next Steps (1–3)

This package contains:

1. Training Harness (modulus/training/harness.py)

   • JAX-first, JIT-friendly step function builder
   • Hyperball metrics extraction
   • Optional LoRA gradient steering hook

2. Parameter-group Hyperball (modulus/optim/groups.py)

   • Path-aware labeler for params (attn vs mlp vs embed vs norm vs bias)
   • optax.multi_transform builder so each group can have distinct Hyperball configs
   • Designed for clean ablations

3. LoRA tangent-steering (modulus/peft/lora.py)

   • Minimal Flax LoRA Dense module
   • orth_lora_grad_jax + apply_lora_grad_hook to orthogonalize LoRA factor gradients
   • Works alongside Hyperball constraints on base weights

Quick start

Install (development):

python -m pip install -e ".[dev]"

Run unit tests:

python -m pytest

Run the grouped Hyperball + LoRA demo:

python -m pip install -e ".[examples]"
python -m modulus.examples.train_grouped_hyperball_lora_demo

Run ablation benchmarks (CSV artifacts):

python scripts/run_benchmarks.py

Run benchmark on a real-world streamed corpus (SlimPajama / MiniPile-style):

python scripts/run_benchmarks.py \
  --data-source hf_http \
  --dataset-name HuggingFaceFW/fineweb \
  --dataset-config sample-10BT \
  --dataset-train-split train \
  --dataset-eval-split train \
  --dataset-tokenizer-backend tiktoken \
  --dataset-tokenizer-name cl100k_base \
  --dataset-eval-holdout-fraction 0.01 \
  --hardware-aware \
  --auto-adjust-max-steps-for-token-target \
  --param-dtype auto \
  --max-tokens-per-step 4096 \
  --max-logits-elements 33554432 \
  --max-attention-elements 8388608 \
  --auto-seq-len-by-memory \
  --auto-disable-distill-for-memory \
  --compile-retry-attempts 3 \
  --compile-heartbeat-sec 30 \
  --telemetry-memory-interval 25 \
  --inference-sampler-interval 100 \
  --inference-sampler-temperature 1.0 \
  --hellaswag-eval-interval 100 \
  --hellaswag-max-examples 128 \
  --train-pool-refresh-interval 250 \
  --auto-token-pool-by-host-ram \
  --host-ram-token-pool-fraction 0.20 \
  --dataset-http-cache-dir artifacts/datasets/hf_http_cache \
  --dataset-http-cache-read \
  --dataset-http-cache-write \
  --dataset-token-cache-dir artifacts/datasets/token_pool_cache \
  --dataset-token-cache-read \
  --dataset-token-cache-write \
  --dataset-token-cache-prime-train-tokens 8388608 \
  --lr 6e-4 \
  --lr-schedule warmup_cosine \
  --lr-warmup-steps 500 \
  --lr-min-ratio 0.10

Note: for hf_http, keep --dataset-rows-page-size at <=100 (HF API limit). For higher request budgets, set HF_TOKEN in the environment and the runner will send Authorization: Bearer ... to dataset-server. The default real-data tokenizer backend is tiktoken (cl100k_base by default). Alternative backends: --dataset-tokenizer-backend hash (legacy regex/hash) and --dataset-tokenizer-backend hf_auto --dataset-tokenizer-name <hf_tokenizer_id>. External token IDs are projected into model vocab with --dataset-token-id-projection table (default, stable ID table + UNK overflow). Legacy modulo projection remains available via --dataset-token-id-projection mod. When train and eval use the same split, set --dataset-eval-holdout-fraction 0.01 to keep validation isolated as a deterministic 1% holdout. If rate-limited, increase --dataset-http-max-retries and --dataset-http-min-interval-sec. Use --log-interval (for example 10) to print rich live progress lines. For long runs, increase --step-record-interval (for example 10 or 25) to reduce memory and CSV size while preserving eval snapshots. TPU observability knobs: --compile-heartbeat-sec emits periodic heartbeat logs during first-step XLA compile, and --telemetry-memory-interval samples host/device memory proxies into benchmark_steps.csv. Progress-quality hooks: --inference-sampler-interval writes periodic generation samples to inference_samples.jsonl (temperature defaults to 1.0), and --hellaswag-eval-interval tracks HellaSwag accuracy over training. For long real-data runs, set --train-pool-refresh-interval (for example 250) so training does not overfit a fixed token pool. Optional profiler capture: --profile-trace --profile-trace-dir <path> and/or --profile-server-port <port> for TensorBoard profiler attachment. Hardware-aware mode is on by default and can downshift batch_size or token_pool_batches when requested settings exceed device/host limits. For hf_stream mode, set HF_HOME / HF_DATASETS_CACHE to a persistent path to reuse downloaded shards across reruns. Token pool cache is shape-agnostic (batch/seq changes reuse the same cached pool when dataset/tokenizer/partition settings match). Use --dataset-token-cache-prime-train-tokens to prefill a larger train pool once and amortize later runs.

To pre-stage data only (no model init/JIT), run:

python scripts/run_benchmarks.py \
  --data-source hf_http \
  --dataset-name HuggingFaceFW/fineweb \
  --dataset-config sample-10BT \
  --dataset-tokenizer-backend tiktoken \
  --dataset-tokenizer-name cl100k_base \
  --dataset-eval-holdout-fraction 0.01 \
  --prepare-data-only \
  --dataset-token-cache-prime-train-tokens 8388608

If a dataset is unavailable in your environment, switch to another public stream (for example --dataset-name cerebras/SlimPajama-627B) or log in with huggingface-cli login for gated datasets.

Colab TPU substantial baseline (single config, 1B-token budget):

python scripts/run_benchmarks.py \
  --data-source hf_http \
  --dataset-name HuggingFaceFW/fineweb \
  --dataset-config sample-10BT \
  --dataset-tokenizer-backend tiktoken \
  --dataset-tokenizer-name cl100k_base \
  --dataset-eval-holdout-fraction 0.01 \
  --configs baseline \
  --target-train-tokens 1000000000 \
  --steps 1000 \
  --max-steps 90000 \
  --lr 6e-4 \
  --lr-schedule warmup_cosine \
  --lr-warmup-steps 2000 \
  --lr-min-ratio 0.10 \
  --lr-total-steps 90000 \
  --compile-heartbeat-sec 30 \
  --telemetry-memory-interval 25 \
  --inference-sampler-interval 500 \
  --inference-sampler-temperature 1.0 \
  --hellaswag-eval-interval 500 \
  --hellaswag-max-examples 256 \
  --train-pool-refresh-interval 250

Colab TPU v6e1 100M-token all-ablation run (~140M model, 5 configs):

python scripts/run_benchmarks.py \
  --data-source hf_http \
  --dataset-name HuggingFaceFW/fineweb \
  --dataset-config sample-10BT \
  --dataset-tokenizer-backend tiktoken \
  --dataset-tokenizer-name cl100k_base \
  --dataset-token-id-projection table \
  --dataset-eval-holdout-fraction 0.01 \
  --configs baseline,lora_hook_only,hyperball_ungrouped,hyperball_grouped,hyperball_grouped_lora \
  --target-train-tokens 100000000 \
  --steps 2000 \
  --max-steps 30000 \
  --batch-size 8 \
  --seq-len 1024 \
  --width 768 \
  --num-layers 12 \
  --num-heads 12 \
  --vocab-size 32768 \
  --hardware-aware \
  --max-tokens-per-step 12288 \
  --max-logits-elements 268435456 \
  --max-attention-elements 100663296 \
  --eval-interval 500 \
  --eval-batches 4 \
  --inference-sampler-interval 500 \
  --inference-sampler-temperature 1.0 \
  --hellaswag-eval-interval 500 \
  --hellaswag-max-examples 256 \
  --train-pool-refresh-interval 250 \
  --lr 3e-4 \
  --lr-schedule warmup_cosine \
  --lr-warmup-steps 1000 \
  --lr-min-ratio 0.10 \
  --lr-total-steps 30000

Build benchmark report from CSV artifacts:

python -m pip install -e ".[report]"
python scripts/build_benchmark_report.py --with-plots

Ablation switches

• Turn off Hyperball per group by setting that group's hb_kwargs_by_group[label] = {}.
• Turn off LoRA gradient steering via use_lora=False in the config (or use_lora_grad_hook=False).
• Switch Hyperball granularity: "leaf" vs "row" vs "col" vs "channel".
• Try "ball" mode with radial_decay and ball_norm_clamp.

IP and compliance

• License (MIT): LICENSE
• Provenance record: docs/ip/provenance.md
• AI origin evidence: docs/ip/ai_origin_evidence.md
• Diligence signoff: docs/ip/SIGNOFF.md
• Third-party notices inventory: THIRD_PARTY_NOTICES.md
• Diligence report: IP_DUE_DILIGENCE_REPORT_2026-03-04.md
• Contribution policy: CONTRIBUTING.md
• Security policy: SECURITY.md

Engineering handoff

• GCT handoff plan: docs/engineering/ENGINEERING_HANDOFF_GCT_2026-03-04.md
• Install flow: docs/engineering/INSTALL.md
• Lint phase plan: docs/engineering/LINT_PHASE_PLAN.md
• Integration guide: docs/integration_guide.md
• Documentation hub: docs/DOCUMENTATION.md
• Colab pedagogy notebook: notebooks/MODULUS_Pedagogical_Walkthrough.ipynb
• Artifacts layout: artifacts/README.md
• CI workflow: .github/workflows/ci.yml

Presets

• LLaMA/HF decoder presets: modulus.optim.presets