Skip to content

[Feat] Add native Tensor Parallelism support for HF backend#3692

Open
YangKai0616 wants to merge 2 commits intoEleutherAI:mainfrom
YangKai0616:hf-tp
Open

[Feat] Add native Tensor Parallelism support for HF backend#3692
YangKai0616 wants to merge 2 commits intoEleutherAI:mainfrom
YangKai0616:hf-tp

Conversation

@YangKai0616
Copy link
Copy Markdown

@YangKai0616 YangKai0616 commented Apr 9, 2026

Using TP to accelerate lm_eval for HF.

Usage:
With torchrun:

torchrun --nproc-per-node=4 -m lm_eval \
    --model hf \
    --model_args pretrained=google/gemma-4-31B-it,max_length=4096,add_bos_token=True,tp_plan=auto \
    --tasks gsm8k \
    --num_fewshot 5 \
    --batch_size 16 \
    --output_path ./GSM8K/31B

Or with accelerate launch:

accelerate launch --num_processes=4 -m lm_eval \
    --model hf \
    --model_args pretrained=google/gemma-4-31B-it,max_length=4096,add_bos_token=True,tp_plan=auto \
    --tasks gsm8k \
    --num_fewshot 5 \
    --batch_size 16 \
    --output_path ./GSM8K/31B

@YangKai0616 YangKai0616 requested a review from 0xSMT as a code owner April 9, 2026 06:43
@CLAassistant
Copy link
Copy Markdown

CLAassistant commented Apr 9, 2026

CLA assistant check
All committers have signed the CLA.

@YangKai0616
Copy link
Copy Markdown
Author

Hi @baberabb ,please help review this PR, thanks!
Cc @sywangyi

@YangKai0616
Copy link
Copy Markdown
Author

Hi @baberabb , gentle ping on this PR. Let me know if you need any changes or further information from my side. Thanks!

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