HELLoRA replication

An unofficial, single-GPU replication of HELLoRA on the pinned allenai/OLMoE-1B-7B-0125 backbone and GSM8K.

Version v0.1.0 is a research preview. It implements layer-level expert profiling, selective HELLoRA injection, matched full LoRA, adapter-only checkpointing, and held-out evaluation. It is intentionally transparent about the extraction sensitivity of the base-versus-adapter comparison and the single-seed uncertainty of the adapter comparison.

Released artifacts

Artifact	Contents
HELLoRA checkpoint	Top-8 experts per layer, adapter weights, topology, and metrics
Full LoRA checkpoint	All-expert adapter weights, topology, and metrics
Evaluation archive	All 3,957 base/HELLoRA/LoRA responses, scores, reports, and checksums

The adapters use this project's strict custom safetensors format, not the PEFT checkpoint format. Load them with the injection and checkpointing code in this repository.

Result summary

All methods use the same pinned 1,319-example GSM8K test split, prompt, greedy decoding, 512-token cap, and numeric answer extractor.

Method	Parameters	Format-tolerant	Strict ####	Common final number
Pretrained base	0	46.47%	0.00%	1.97%
Full LoRA, seed 42	311,459,840	41.17%	41.17%	41.02%
HELLoRA, seed 42	47,218,688	39.58%	39.58%	39.58%

The format-tolerant extractor uses ####, then the first explicit answer phrase, then the final number. It recovers early base-model answers before the base continues into unrelated examples; the adapters almost always emit #### and stop. Consequently, 46.47% versus the adapter scores is not a parser-neutral ranking. The strict shared marker regex scores the base at 0.00%, while a shared final-number regex scores it at 1.97%.

Full LoRA leads HELLoRA by 1.59 points, but the paired result is not statistically decisive (p = 0.286; paired-bootstrap 95% interval -1.21 to +4.40 points). HELLoRA uses 84.84% fewer adapter parameters and produces a 6.60-times smaller checkpoint. See the paired report and base comparison.

The LoRA-versus-HELLoRA comparison is less sensitive because both adapters use the same output format. Full LoRA leads by 1.59 points under format-tolerant scoring and 1.44 points under common-final-number scoring; neither paired gap is statistically decisive. These numbers are not a direct reproduction of the paper's table: this project uses a newer OLMoE checkpoint, one seed, and explicit implementation choices where the paper does not fully specify the protocol.

Implemented scope

• pinned model and dataset revisions with local-only loading support;
• hard per-layer router-assignment profiling with padding exclusion;
• deterministic layer-hot and layer-cold expert selection;
• LoRA on attention projections, routers, and selected or all experts;
• LoRA-only trainability enforcement and exact no-op initialization;
• single-GPU warm-up and stage-two training with memory accounting;
• adapter-only safetensors checkpoints with SHA-256 and exact reload parity;
• held-out GSM8K generation, finalized answer extraction, rescoring, and paired comparison;
• 39 unit and integration tests.

Mixtral, DeepSeekMoE, HELLoRI, multi-GPU execution, and multi-seed confirmation remain out of scope for this research preview. The detailed project decisions are recorded in the implementation plan.

Requirements

The completed training runs used:

• one NVIDIA RTX PRO 6000 Blackwell Workstation Edition with 94.97 GiB VRAM;
• Python 3.12.13;
• PyTorch 2.12.0 with CUDA 13.0;
• Transformers 4.53.3 and Datasets 3.6.0.

The repository does not claim that the declared training schedule fits a smaller GPU. The base model and GSM8K revisions are downloaded separately from Hugging Face.

Installation

For normal development:

python -m pip install -e '.[olmoe,dev]'
pytest -q

For closer reconstruction of the completed container, use Python 3.12 with a compatible CUDA 13 driver and install the captured environment first:

python -m pip install -r requirements-freeze.txt
python -m pip install --no-deps -e .
pytest -q

requirements-freeze.txt includes development and notebook packages and is an environment snapshot rather than the minimal supported dependency set.

Download released adapters

hf download \
  toolazyhhh123/hellora-olmoe-1b-7b-gsm8k-seed42 \
  --local-dir checkpoints/hellora-olmoe-gsm8k-seed42

hf download \
  toolazyhhh123/lora-olmoe-1b-7b-gsm8k-seed42 \
  --local-dir checkpoints/lora-olmoe-gsm8k-seed42

The base model is pinned to revision 9b0c1aa87e34a20052389dce1f0cf01da783f654, and GSM8K is pinned to revision 740312add88f781978c0658806c59bc2815b9866.

Reproduce the pipeline

Profile the 10% warm-up sample:

hellora-warmup \
  --output-dir reports/warmup-olmoe-gsm8k-seed42 \
  --cache-dir /path/to/huggingface/cache

Train selective HELLoRA:

hellora-finetune \
  --method hellora \
  --hot-experts reports/warmup-olmoe-gsm8k-seed42/hot_experts.json \
  --checkpoint-dir checkpoints/hellora-olmoe-gsm8k-seed42 \
  --report-dir reports/finetune-hellora-olmoe-gsm8k-seed42 \
  --cache-dir /path/to/huggingface/cache \
  --micro-batch-size 64 \
  --gradient-accumulation-steps 2

Train the matched full-LoRA baseline by replacing the method and output paths:

hellora-finetune \
  --method lora \
  --checkpoint-dir checkpoints/lora-olmoe-gsm8k-seed42 \
  --report-dir reports/finetune-lora-olmoe-gsm8k-seed42 \
  --cache-dir /path/to/huggingface/cache \
  --micro-batch-size 64 \
  --gradient-accumulation-steps 2

Evaluate either adapter with the same generation protocol:

hellora-evaluate-gsm8k \
  --method hellora \
  --checkpoint-dir checkpoints/hellora-olmoe-gsm8k-seed42 \
  --output-dir artifacts/evaluations/gsm8k-hellora-seed42 \
  --cache-dir /path/to/huggingface/cache \
  --batch-size 32 \
  --max-new-tokens 512

The OLMoE config declares float32, so the loader explicitly forces BF16. Do not replace that behavior silently when reproducing the memory measurements.

Contributing and security

See CONTRIBUTING.md for development guidance. Report security issues privately as described in SECURITY.md.

Citation

Use CITATION.cff to cite this software. Please also cite the original HELLoRA paper when this implementation or its artifacts support research derived from that method.

License and attribution

The implementation is licensed under Apache-2.0. The converted paper in HELLoRA-paper/ is separate third-party material licensed CC BY 4.0 and is not covered by the software license. See NOTICE and HELLoRA-paper/README.md for attribution and change disclosure.