Commit f950025
committed
feat: LLM backbone encoder for JEPA training with local model selection
Add support for using open-weights LLMs (Llama, Qwen, Gemma, Mistral, Phi)
as frozen feature extractors for the JEPA training pipeline. The LLM's
hidden states provide rich semantic embeddings; only a small projection
head (~2-8M params) is trained and distributed via FedAvg. The LLM
weights never leave the node.
New files:
- nodes/common/local_models.py: Model catalog (11 models), host hardware
probe (RAM, VRAM, disk, MPS), compatibility filtering, recommendation
- nodes/common/llm_backbone.py: LLMBackboneEncoder (frozen LLM + trainable
projection head), BackboneJEPATrainer (JEPA training with LLM features),
learned layer mixing, EMA target encoder
Integration:
- ml.py: New backbone training path in train_vljepa_on_task() when
backbone_model_id is set in task_spec
- text_data.py: Include raw_texts in batch dict for LLM re-tokenization
- training_feed.py: Pass backbone_model_id through to task spec
- config.py: Add backbone_model_id to ModelConfig
- service.py: Wire backbone_model_id from config to training feed,
fix governance variable reference bug
- ws_server.py: Add local_models and set_local_model WS endpoints
- pyproject.toml: Add transformers + accelerate to [network] deps1 parent 88580b6 commit f950025
9 files changed
Lines changed: 1380 additions & 3 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
900 | 900 | | |
901 | 901 | | |
902 | 902 | | |
| 903 | + | |
| 904 | + | |
| 905 | + | |
| 906 | + | |
| 907 | + | |
| 908 | + | |
| 909 | + | |
| 910 | + | |
| 911 | + | |
| 912 | + | |
| 913 | + | |
| 914 | + | |
| 915 | + | |
| 916 | + | |
| 917 | + | |
| 918 | + | |
| 919 | + | |
| 920 | + | |
| 921 | + | |
| 922 | + | |
| 923 | + | |
| 924 | + | |
| 925 | + | |
| 926 | + | |
| 927 | + | |
| 928 | + | |
| 929 | + | |
| 930 | + | |
| 931 | + | |
| 932 | + | |
| 933 | + | |
| 934 | + | |
| 935 | + | |
903 | 936 | | |
904 | 937 | | |
905 | 938 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
50 | 50 | | |
51 | 51 | | |
52 | 52 | | |
53 | | - | |
| 53 | + | |
54 | 54 | | |
55 | 55 | | |
56 | 56 | | |
57 | 57 | | |
58 | 58 | | |
59 | 59 | | |
60 | 60 | | |
| 61 | + | |
| 62 | + | |
61 | 63 | | |
62 | 64 | | |
63 | 65 | | |
| |||
0 commit comments