Skip to content

Support custom-code backbones (Nandi, Lumma) and training on from decisionformer - #11

Open
Rohit7sali wants to merge 1 commit into
mainfrom
custom-model-support
Open

Rohit7sali wants to merge 1 commit into
mainfrom
custom-model-support

Conversation

@Rohit7sali

Copy link
Copy Markdown

Support custom-code backbones (Nandi, Lumma) and training on from decision checkpoints

Adds inference and fine-tuning for FrontiersMind's models:

Model Inference Fine-tuning
Lumma-fev-0.1B-v2 (custom Nandi code) yes, trust_remote_code=True / --trust-remote-code trains on from it
Lumma-fev-4B-v2, Lumma-fev-27B-v2 (Qwen3.5) yes (already loaded; now documented) trains on from them (LoRA / QLoRA)
Nandi-Mini-150M, Lumma-0.6B-Base (custom code) n/a yes
Other Hub decoders with their own modeling code yes yes, model.trust_remote_code: true

Changes

  • trust_remote_code end to end: new --trust-remote-code flag on eval / predict / serve; LoRA merge passes the flag along; exports copy the model's .py files and auto_map, so checkpoints load on their own. Loading without the flag fails with a message that names it.
  • Nandi/Lumma plain prompt (data.prompt.plain: auto | decisionformer | lumma): <bos> + prompt + \n\nAnswer:, tokenized with the fast tokenizer so the custom encode() cannot add a second BOS. auto picks lumma for Nandi/Lumma model types; the choice is saved as plain_prompt in decision_config.json. Other models keep the existing prompt byte for byte.
  • Factorized embeddings: the code readout starts from the LM head through lm_head_proj, so it reproduces the base model's code-token logits at step 0.
  • Training on from a decision checkpoint: when model.backbone has a decision_config.json, its readout seeds the code_readout head, and its prompt, attention mode and codes carry over unless set in the config. Adapter-only checkpoints are refused with a clear message.
  • Answer-code detection uses the fast tokenizer directly (same ids for standard tokenizers; robust to custom encode overrides).
  • Recipes: nandi_mini, lumma_base (LoRA), lumma_fev_continue.
  • Docs: README sections on custom-code backbones and on training on from a decision model.

Tests

  • tests/test_custom_models.py: 18 offline tests on a tiny custom-code model built like Nandi/Lumma (factorized embedding, lm_head_proj, BOS-prepending tokenizer). They cover the exact prompt tokens, readout = LM-head logits, train → export → reload, LoRA merge, the refusal without the flag, training on from a 4B/27B-v2-style and a 0.1B-v2-style checkpoint (identical outputs at step 0), the CLI flag and the recipes.
  • tests/test_hub_models.py: real checkpoints, opt in with DECISIONFORMER_HUB_TESTS=1 (DECISIONFORMER_HUB_LARGE=1 adds 4B/27B-v2 on a GPU). On an RTX 3060: 8 passed, 1 skipped (the lumma-fev package comparison, skipped when the package is not installed; it passed in a local run), 27B not run. Includes Lumma-fev-0.1B-v2 matching an independent re-implementation of the lumma-fev runtime, Nandi / Lumma base fine-tuning, and Lumma-fev-4B-v2 load + LoRA training.
  • CLI smoke runs on GPU: train nandi_mini → eval --trust-remote-code (reproduces the training-time accuracy exactly) → serve --trust-remote-code answers /v1/systemone. train local_smollm2 → eval unchanged for standard models.
  • Full offline suite: no new failures. The 5 gemma2-sdpa failures in tests/test_attention.py also fail on main since transformers 5.21 (Gemma2 no longer accepts SDPA); they will be fixed in a separate PR.

@Rohit7sali Rohit7sali changed the title Support custom-code backbones (Nandi, Lumma) and training on from dec… Support custom-code backbones (Nandi, Lumma) and training on from decisionformer Nov 2, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants