riccardozaccone96/flexrank-llama3.2-1B

Llama 3.2 1B FlexRank

This checkpoint applies FlexRank to Llama 3.2 1B. It contains a single consolidated set of nested low-rank weights that can run at different parameter budgets without loading or training a separate checkpoint for each size.

FlexRank orders the low-rank components of decomposed linear layers by importance. At inference time, reduce_size selects a stored rank profile for the requested budget, enabling a gradual quality-versus-cost trade-off from one model.

License: the FlexRank software is MIT-licensed. These model weights are derived from Llama 3.2 and are distributed under the Llama 3.2 Community License.

Model details

• Base model: meta-llama/Llama-3.2-1B
• Model type: decoder-only autoregressive Transformer with nested low-rank linear layers
• Method: DataSVD initialization, dynamic-programming profile search, and KL-only knowledge consolidation
• Input/output: multilingual text to text
• Languages: English, German, French, Italian, Portuguese, Hindi, Spanish, and Thai are inherited as officially supported languages from Llama 3.2. The FlexRank checkpoint and its reduced-size profiles have currently been evaluated only on English benchmarks; multilingual quality after consolidation has not been measured.
• Maximum context inherited from the base model: 128k tokens
• Training data: FineWeb-Edu sample-10BT
• License: Llama 3.2 Community License

Sources

• FlexRank repository: https://github.com/RickZack/FlexRank
• Project page: https://rickzack.github.io/FlexRank/
• Paper: https://openreview.net/forum?id=DK0kvnNelx
• Usage notebook: https://github.com/RickZack/FlexRank/blob/main/notebooks/flexrank_save_load_prune.ipynb

Intended use

This checkpoint is intended for research on elastic language models, adaptive inference, nested low-rank representations, and deployment under changing memory or compute budgets. It can also be used as a pretrained causal language model or as a starting point for downstream fine-tuning.

This is a base pretrained model, not an instruction-tuned assistant. It should not be expected to follow conversational instructions reliably without additional adaptation. Lower parameter budgets trade language-model quality for reduced computation and should be validated for the target application.

All uses remain subject to the Llama 3.2 Community License and Acceptable Use Policy. The model is not intended for high-stakes decisions or for generating factual, medical, legal, or financial advice without appropriate safeguards and expert review.

Usage

Loading requires the FlexRank runtime package and trust_remote_code=True. The Hub repository includes the custom flexrank_model.py Transformers wrapper, while its low-rank layers, profiles, samplers, and utilities are imported from the flexrank package in the FlexRank source repository.

Install the runtime package directly from the repository before loading the model:

pip install "git+https://github.com/RickZack/FlexRank.git#subdirectory=flexrank"

The loading call uses trust_remote_code=True to execute the custom wrapper. Review that code before using the checkpoint in a security-sensitive environment. reduce_size supports both reversible virtual profile selection and irreversible physical deployment.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from flexrank.samplers.base_sampler import DeployMode

model_id = "riccardozaccone96/flexrank-llama3.2-1B"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.float32,
    device_map="auto",
)

# Reversible virtual selection (the default): activate the stored profile nearest
# to 60% while retaining the full factors. Another profile can be selected later.
model.reduce_size(size_ratio=0.6, deploy_mode=DeployMode.NO)

inputs = tokenizer("FlexRank enables adaptive inference by", return_tensors="pt").to(model.device)
with torch.inference_mode():
    output_ids = model.generate(**inputs, max_new_tokens=40)
print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

# Optional physical deployment after choosing the final budget. This slices the
# stored factors, so the model cannot subsequently be expanded to a larger size.
# model.reduce_size(size_ratio=0.6, deploy_mode=DeployMode.SVD)
# model.save_pretrained("flexrank-llama3.2-1B-svd-60pct")

size_ratio is a target relative FlexRank parameter budget, not a guarantee that the physical checkpoint file or every non-decomposed part of the model shrinks by the same percentage. The model selects the closest stored profile that does not exceed the requested budget.

Deployment mode	Effect	Can select a larger profile later?
DeployMode.NO	Changes only the active ranks; retains the complete stored factors. This is the default and is recommended for comparing sizes.	Yes
DeployMode.SVD	Physically slices the two low-rank factors at the selected ranks. Use this to save a conventional compact checkpoint.	No
DeployMode.GAR	Physically converts supported layers to the more compact gauge-aligned representation; unsupported layers fall back to SVD.	No

Physical deployment mutates the model in place. Reload the full Hub checkpoint before creating a different physically deployed size. GAR deployment should be checked against the equivalent virtual profile in FP32; its gauge-aligned factors can be numerically sensitive under BF16 inference.

Training

The base Llama weights were decomposed with DataSVD using 256 calibration sequences from FineWeb-Edu. A dynamic-programming search constructed ten nested profiles between 10% and 100% of the FlexRank parameter budget. Knowledge consolidation then trained the shared factors while sampling predefined profiles.

Training used only the KL-divergence distillation objective against the original Llama model (ce_loss_w=0, kl_loss_w=1, temperature 1). The run processed approximately 60 billion token presentations by repeating the FineWeb-Edu sample-10BT data.

Hyperparameter	Value
Training steps	120,000
Sequence length	1,024
Global batch size	512
Precision	BF16 mixed precision
Optimizer	AdamW
Peak learning rate	6e-5
Minimum learning rate	6e-6
LR schedule	Cosine with minimum LR
Warmup steps	715
Adam betas	0.9, 0.95
Weight decay	0
Gradient clipping	1.0
Seed	1

The language-model head was excluded from low-rank decomposition. Training did not use a cross-entropy term; it preserved the teacher distribution through knowledge consolidation.

Evaluation

The model was evaluated one-shot on the complete standard evaluation splits using a modified checkout of lm-evaluation-harness version 0.4.9.1 with the FlexRank model adapter. The evaluation used batch size 32.

Here, one-shot means that one in-context demonstration example was provided for each task (num_fewshot=1). These multiple-choice tasks score answer choices and report accuracy; this is not pass@1, which is normally used for sampled-generation tasks where generated candidates are tested for success.

All six tasks report accuracy (acc). All values are percentages; uncertainty is the standard error reported by the harness. The macro average is the unweighted mean across the six tasks and is included for readability, not as a separately registered benchmark.

One-shot results by parameter budget

Parameter budget	ARC-C	ARC-E	HellaSwag	OpenBookQA	PIQA	WinoGrande	Macro avg.
100%	34.73 ± 1.39	68.18 ± 0.96	45.26 ± 0.50	28.80 ± 2.03	74.10 ± 1.02	60.54 ± 1.37	51.94
90%	33.70 ± 1.38	67.89 ± 0.96	45.04 ± 0.50	28.20 ± 2.01	73.23 ± 1.03	60.54 ± 1.37	51.43
80%	33.53 ± 1.38	67.76 ± 0.96	44.64 ± 0.50	27.80 ± 2.01	73.78 ± 1.03	59.51 ± 1.38	51.17
70%	33.62 ± 1.38	66.88 ± 0.97	43.91 ± 0.50	28.80 ± 2.03	73.61 ± 1.03	58.33 ± 1.39	50.86
60%	31.14 ± 1.35	65.40 ± 0.98	42.80 ± 0.49	27.40 ± 2.00	73.18 ± 1.03	57.14 ± 1.39	49.51
50%	30.20 ± 1.34	64.52 ± 0.98	41.15 ± 0.49	27.00 ± 1.99	71.44 ± 1.05	57.62 ± 1.39	48.66
40%	28.41 ± 1.32	63.01 ± 0.99	39.28 ± 0.49	27.00 ± 1.99	69.86 ± 1.07	55.49 ± 1.40	47.18
30%	27.22 ± 1.30	59.76 ± 1.01	36.81 ± 0.48	27.20 ± 1.99	68.06 ± 1.09	53.99 ± 1.40	45.51
20%	22.95 ± 1.23	54.50 ± 1.02	32.49 ± 0.47	20.60 ± 1.81	65.23 ± 1.11	53.04 ± 1.40	41.47
10%	18.00 ± 1.12	31.19 ± 0.95	26.25 ± 0.44	12.60 ± 1.49	56.86 ± 1.16	49.80 ± 1.41	32.45

Limitations and risks

• The model inherits the biases, risks, multilingual performance differences, and limitations of Llama 3.2 1B and of the FineWeb-Edu data used for consolidation.
• This checkpoint is not instruction tuned and can produce inaccurate, biased, unsafe, or offensive text.
• Benchmark performance generally decreases with the active parameter budget, although individual finite-sample scores need not be strictly monotonic.
• The reported benchmarks cover English multiple-choice reasoning and commonsense tasks; they do not characterize safety, factuality, long-context behavior, generation quality, or performance across all supported languages.
• Physically deployed GAR factors can be numerically sensitive under reduced-precision inference. Use FP32 verification or SVD deployment when numerical equivalence is required.
• Loading executes custom repository code through trust_remote_code=True; review the code before using it in security-sensitive environments.

Environmental impact

Energy use, hardware model, and carbon emissions for this consolidation run were not recorded, so an emissions estimate is not provided.

Citation

@inproceedings{
  zaccone2026flexrank,
  title={FlexRank: Nested Low-Rank Knowledge Decomposition for Adaptive Model Deployment},
  author={Zaccone, Riccardo and Laskaridis, Stefanos and Ciccone, Marco and Horvath, Samuel},
  booktitle={Forty-third International Conference on Machine Learning},
  year={2026},
  url={https://openreview.net/forum?id=DK0kvnNelx}
}

Contact

For questions and issues, use the FlexRank GitHub repository.