Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

feat: add tensor-parallel entropy loss benchmark

Đang mở
  • Có gì không ổn?
#500 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
25/100
Loại issue
Tính năng
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Ít trao đổi
Công nghệ
python
Lĩnh vực
devtools, performance

Hướng nghiên cứu

Start with tools/benchmark/bench_entropy_loss.py to see the existing single-GPU variants, then read tools/benchmark/runner.py and main.py to understand the current single-process flow. The issue leaves open whether to add a new tools/benchmark/run_tp.py (Option A) or extend main.py with a --tp flag (Option B), so settle that with a maintainer first. Done means the three TP variants run across the five (vocab, tp_size) shapes with synchronized timing and OOM reported for pytorch_gather.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

Motivation

The Triton entropy loss kernel supports tensor-parallel training natively via its group parameter: each rank holds vocab / tp_size logits, all-reduces two scalars per row (max and sum-exp), then computes loss and grad locally. This avoids materializing the full logits on any rank.

The PyTorch alternative requires an all-gather of the full logit tensor first — O(tokens × vocab) communication. At realistic scale this is prohibitive:

  • Llama 3.1 405B: vocab=128K, tokens=8K, TP=8 → 16 GB all-gather per step

The Triton path is not just faster; it is the only feasible approach at large vocab × TP. The current single-GPU benchmark does not demonstrate this. This issue tracks building a multi-GPU benchmark that makes it concrete.

Variants to benchmark

Three qualitatively different approaches:

  1. triton_tp — existing triton_entropy_loss_forward_backward(..., group=group). Shards vocab across ranks, all-reduces two scalars per row. O(tokens) communication.

  2. pytorch_tp_manual — same algorithm in PyTorch without Triton: local_max = logits.max(-1) → all_reduce(MAX) → local_sum = (logits − max).exp().sum(-1) → all_reduce(SUM) → loss. Tests whether Triton fusion still wins when both paths use the same O(tokens) communication pattern.

  3. pytorch_gather — all-gather logits to full vocab on each rank, then F.cross_entropy. O(tokens × vocab) communication. Included as a reference to show where the naive approach becomes infeasible; expected to OOM at large vocab × TP.

Shapes

Fix tokens=4096. Sweep (vocab, tp_size) pairs:

vocab tp_size shard / rank
32768 2 16384
32768 4 8192
65536 4 16384
131072 4 32768
131072 8 16384

Infrastructure changes

Multi-process runner

The current runner.py is single-process. Two options:

Option A — new tools/benchmark/run_tp.py entry point using torch.multiprocessing.spawn(worker, nprocs=tp_size). Each worker: initializes dist.init_process_group, creates a TP process group, runs benchmark variants with group=group, rank 0 collects and prints results.

Option B — extend __main__.py to detect a --tp N flag and re-launch itself via torchrun --nproc_per_node=N.

Option A is simpler to implement. Option B integrates more cleanly with the existing CLI.

Timing
  • dist.barrier() + torch.cuda.synchronize() before and after each timed region so all ranks agree on wall time.
  • Report max latency across ranks (rank 0 collects via dist.all_reduce(MAX)).
  • Communication time is included automatically since the all-reduce is inside the kernel call.
OOM guard for pytorch_gather

Wrap in try/except and report OOM in the table instead of a time. This makes the table show exactly where the naive approach becomes infeasible.

Expected outcome

At vocab=131072, TP=8:

variant result
triton_tp fast (~O(tokens) communication)
pytorch_tp_manual slightly slower (no kernel fusion, same communication)
pytorch_gather OOM or ~10–20× slower (16 GB all-gather dominates)

Files

  • New: tools/benchmark/run_tp.py (or extend __main__.py)
  • Modified: tools/benchmark/bench_entropy_loss.py — add TP variants alongside existing single-GPU variants
Ngôn ngữ chính
Python
Star
337
Fork
52
Không có pull request nào được merge trong 30 ngày

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của ServiceNow/Fast-LLM

Tất cả issue của ServiceNow/Fast-LLM

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.