Skip to content

feat: add tensor parallelism awareness for multi-GPU model fitting#277

Merged
AlexsJones merged 2 commits into
AlexsJones:mainfrom
michaeljabbour:feat/tensor-parallelism
Mar 23, 2026
Merged

feat: add tensor parallelism awareness for multi-GPU model fitting#277
AlexsJones merged 2 commits into
AlexsJones:mainfrom
michaeljabbour:feat/tensor-parallelism

Conversation

@michaeljabbour

Copy link
Copy Markdown
Contributor

Summary

  • supports_tp(tp_size) and valid_tp_sizes() methods on LlmModel for checking TP compatibility
  • Heuristic head-count inference for Qwen, Llama, DeepSeek, Mistral, Gemma, Phi, MiniMax model families
  • RunMode::TensorParallel with 0.9× speed penalty in fit analysis
  • TUI: TP compatibility filter (Shift-T) cycles All → TP=2 → TP=3 → TP=4

Files changed (9)

File Change
llmfit-core/src/models.rs TP fields, methods, head-count inference, tests (+269)
llmfit-core/src/fit.rs TensorParallel run mode, cluster-aware fit path (+54)
llmfit-core/src/hardware.rs cluster_mode, cluster_node_count fields (+10)
llmfit-core/src/plan.rs Test constructor updates (+4)
llmfit-tui/src/tui_app.rs TpFilter enum, cycling, filter integration (+47)
llmfit-tui/src/tui_events.rs Shift-T keybinding (+3)
llmfit-tui/src/tui_ui.rs TP filter widget in search bar (+20)
llmfit-tui/src/serve_api.rs TensorParallel in run_mode_code (+1)
llmfit-tui/src/main.rs Test constructor updates (+6)

Test plan

  • All existing tests pass (245 passed, 1 pre-existing failure unrelated to this PR)
  • New TP unit tests: supports_tp, valid_tp_sizes, infer_heads_from_name for multiple model families
  • Manual: launch TUI, press Shift-T to cycle TP filter, verify model list updates

Split from #181 — see closing comment for context on the split.

🤖 Generated with Claude Code

- supports_tp() and valid_tp_sizes() methods on LlmModel
- Heuristic head count inference for Qwen, Llama, DeepSeek, Mistral,
  Gemma, Phi, MiniMax families
- TensorParallel run mode with 0.9× speed penalty in fit analysis
- TUI: TP compatibility filter (Shift-T) cycles All/TP=2/TP=3/TP=4
- Unit tests for TP logic and head count inference

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@AlexsJones

Copy link
Copy Markdown
Owner

what an awesome PR @michaeljabbour

@AlexsJones AlexsJones merged commit e8a6001 into AlexsJones:main Mar 23, 2026
6 checks passed
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