Các thủ thuật từ OpenAI gpt-oss mà BẠN 🫵 có thể sử dụng với transformers
- 17 min read
Thủ thuật từ OpenAI gpt-oss MỌI NGƯỜI 🫵 có thể sử dụng với transformers
OpenAI gần đây đã phát hành bộ mô hình GPT-OSS. Các mô hình này có một số kỹ thuật mới như lượng tử hóa MXFP4, kernel hiệu quả, định dạng chat hoàn toàn mới và hơn thế nữa. Để cho phép phát hành gpt-oss thông qua transformers, chúng tôi đã nâng cấp đáng kể thư viện. Các bản cập nhật giúp việc tải, chạy và tinh chỉnh các mô hình trở nên rất hiệu quả.
Trong bài đăng trên blog này, chúng ta sẽ nói về tất cả các nâng cấp một cách chuyên sâu và cách chúng trở thành một phần của bộ công cụ transformers để các mô hình khác (hiện tại và tương lai) có thể hưởng lợi từ chúng. Cung cấp các triển khai rõ ràng về các phương pháp mới trong transformers cũng cho phép cộng đồng nhanh chóng hiểu và áp dụng chúng. Các framework như MLX, llama.cpp hoặc vLLM có thể sử dụng mã transformers làm tài liệu tham khảo để xây dựng các triển khai của riêng họ.
Đối với bản phát hành này, chúng tôi đã làm việc trên:
- Kernel không cần build, có thể tải xuống từ Hub
- Lượng tử hóa MXFP4
- Song song hóa Tensor
- Song song hóa Expert
- Lớp Cửa sổ Trượt Động & Bộ nhớ Cache
- Batching liên tục & Chú ý theo Trang
- Tải các mô hình lớn hơn nhanh hơn
Phần hay nhất: Hầu hết các tính năng này sẽ hoạt động trên tất cả các mô hình chính trong
transformers!
Kernel không cần build, có thể tải xuống từ Hub
Kernel là một chương trình chuyên dụng, nhỏ gọn chạy trên các bộ tăng tốc để thực hiện các tác vụ như nhân ma trận, kích hoạt hoặc chuẩn hóa. Trong PyTorch eager, các hoạt động kích hoạt các kernel riêng lẻ một cách tuần tự, điều này rất đơn giản nhưng có thể phát sinh thêm các chuyển bộ nhớ và chi phí khởi chạy. torch.compile của PyTorch 2.0 với các backend như TorchInductor giải quyết vấn đề này bằng cách tự động hợp nhất và tối ưu hóa các kernel, mang lại hiệu suất tăng từ 2–10×.
Ngoài ra, cộng đồng đã tạo ra các kernel tùy chỉnh cho các kết hợp hoạt động thường xuyên, không chỉ các op PyTorch riêng lẻ như matmul. Ví dụ: Flash Attention được tạo ra để tối ưu hóa khối attention quan trọng xác định kiến trúc transformers và có mặt trong nhiều mô hình, bao gồm hầu hết LLM. Bằng cách kết hợp cẩn thận tất cả các hoạt động attention bên trong một kernel duy nhất, các chuyển bộ nhớ được giảm thiểu, việc sử dụng bộ nhớ được giảm và có thể đạt được tăng tốc.
Vấn đề là tất cả các kernel khác nhau này đều có sẵn trong các thư viện riêng biệt, điều này tạo ra sự phình to về dependency nếu chúng được thêm vào thư viện transformers. Hơn nữa, các kernel này không chỉ là mã Python, chúng bao gồm mã cuda cấp thấp, được gắn lại với C++ và được hiển thị thông qua một lớp Python. Điều này có nghĩa là chúng phải được biên dịch trong hệ thống đích, do đó yêu cầu bất kỳ hệ thống build nào được yêu cầu bởi mỗi thư viện kernel.
Gói kernels giải quyết vấn đề này bằng cách tải xuống các binary được build sẵn của các kernel được hỗ trợ từ Hub. Bạn chỉ cần chỉ ra kernel bạn muốn sử dụng và kernels sẽ tìm kiếm một phiên bản tương thích với hệ thống của bạn và tải xuống khi sử dụng lần đầu.
Kernel tùy chỉnh cho GPT-OSS
GPT-OSS, một mô hình Mixture of Experts (MoE), là một người dùng lớn của Kernel từ Hub. Nó tận dụng một số kernel tùy chỉnh:
- Liger RMSNorm, được sử dụng làm
@use_kernel_forward_from_hub("RMSNorm") - Kernel Megablocks MoE:
@use_kernel_forward_from_hub("MegaBlocksMoeMLP") - Flash Attention 3 với hỗ trợ cho attention sinks.
- Kernel triton MXFP4 (được đề cập sau)
Hãy xem hai kernel đầu tiên.
Đằng sau hậu trường, các decorator (1 và 2) chỉ đơn giản là trỏ đến các kernel do cộng đồng đóng góp. Ví dụ: RMSNorm đến từ liger_kernels, trong khi kernel MegaBlocksMoeMLP đến từ megablocks. Tùy thuộc vào thiết bị của bạn (CUDA hoặc ROCm) và bạn đang huấn luyện hay chạy suy luận, kernel phù hợp sẽ được kéo vào tự động.
Thiết kế này vừa cụ thể vừa tổng quát: các kernel RMSNorm liger đã được sử dụng lại trên nhiều mô hình và kernel MoE cũng có thể được áp dụng cho MoE trong tương lai.
Vì kernels kéo mã từ Hub, bạn phải chọn tham gia tính năng này bằng cách chuyển use_kernels=True trong khởi tạo mô hình của bạn, như được hiển thị bên dưới. Chúng tôi bật ghi log INFO trong ví dụ để bạn có thể dễ dàng xác minh rằng các kernel có thể tải xuống đang được sử dụng.
Các kernel này không tương thích với
mxfp4, vì vậy suy luận sẽ xảy ra trongbfloat16nếu bạn sử dụng chúng. Vui lòng benchmark hệ thống của bạn để có sự kết hợp tốt nhất về bộ nhớ và thông lượng phù hợp với dự án của bạn!
python from transformers import AutoTokenizer, AutoModelForCausalLM
import logging logging.basicConfig(level=logging.INFO)
model_id = “openai/gpt-oss-20b” tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained( model_id, dtype=“auto”, device_map=“auto”, use_kernels=True, )
Chạy một thế hệ nhanh chóng tạo ra các thông báo log như
shell
INFO:root:Using layer LigerRMSNorm from repo kernels-community/liger_kernels
INFO:root:Using layer MegaBlocksMoeMLP from repo kernels-community/megablocks
Hình 1 cho thấy rằng, trong hệ thống chúng tôi đã thử nghiệm, các kernel này hoạt động tốt nhất cho các batch size lớn hơn. Chúng tôi luôn khuyên bạn nên benchmark bất kỳ thay đổi nào liên quan đến hiệu suất càng gần với điều kiện sản xuất của bạn càng tốt.
Bạn có thể khám phá và chơi với script benchmark tại đây
Flash Attention 3
Các mô hình OpenAI gpt-oss sử dụng attention sinks, giúp cải thiện chất lượng và tạo điều kiện thuận lợi cho việc sử dụng các ngữ cảnh dài hơn. Nhóm vLLM đã thêm tính năng này vào phiên bản Flash Attention mới nhất (Flash Attention 3) và kernel tùy chỉnh kết quả có sẵn trên Hub. Hiện tại, kernel này tương thích với kiến trúc Hopper. Nếu bạn có một kiến trúc, đây là cách để bật nó:
diff model = AutoModelForCausalLM.from_pretrained( model_id, dtype=“auto”, device_map=“auto”,
-
Flash Attention with Sinks
- attn_implementation=“kernels-community/vllm-flash-attn3”, )
Lượng tử hóa MXFP4
Các mô hình ngôn ngữ lớn rất tốn bộ nhớ. Lượng tử hóa làm giảm footprint bộ nhớ bằng cách lưu trữ weights (và đôi khi là activations) ở các định dạng có độ chính xác thấp hơn. Để tham khảo, FP32 sử dụng 32 bit cho mỗi số và BF16 sử dụng 16. Bằng cách giảm độ rộng bit, chúng ta đánh đổi một số độ chính xác để có các mô hình nhỏ hơn và di chuyển bộ nhớ nhanh hơn.
Nếu bạn muốn có một primer trực quan về các trade-off lượng tử hóa, bài viết của Maarten Grootendorst là tuyệt vời: Hướng dẫn trực quan về lượng tử hóa.
MXFP4 là gì
MXFP4 là một định dạng dấu phẩy động 4-bit với bố cục E2M1: 1 bit dấu, 2 bit số mũ và 1 bit phần định trị, như được hiển thị trong Hình 2. Bản thân E2M1 rất thô. MXFP4 bù đắp bằng scaling theo khối:
- Các vector được nhóm thành các khối gồm 32 phần tử.
- Mỗi khối lưu trữ một scale được chia sẻ khôi phục phạm vi động khi dequantizing.
- Bên trong mỗi khối, các giá trị 4-bit đại diện cho các số liên quan đến scale đó.
Lược đồ theo khối này cho phép MXFP4 giữ phạm vi trong khi sử dụng rất ít bit. Trong thực tế, GPT-OSS 20B phù hợp với khoảng 16 GB VRAM và GPT-OSS 120B phù hợp với khoảng 80 GB khi MXFP4 hoạt động, đó là sự khác biệt giữa “không thể tải” và “có thể chạy trên một GPU duy nhất”. Vấn đề là các phép nhân ma trận hiện phải tôn trọng các scale khối. Làm điều này một cách hiệu quả ở quy mô lớn đòi hỏi các kernel chuyên dụng.
MXFP4 trong transformers
transformers hiện bao gồm hỗ trợ gốc cho MXFP4, tận dụng các kernel triton (MXFP4) được tối ưu hóa để tăng cường hiệu suất. Điều này xây dựng trên phân phối kernel do cộng đồng điều khiển được thảo luận trước đó, sử dụng các kernel được biên dịch trước từ Hub để đơn giản hóa việc triển khai.
Chi tiết triển khai chính:
- Logic quantizer: Được tìm thấy trong tệp quantizer MXFP4, điều này xử lý quá trình lượng tử hóa cốt lõi cho MXFP4.
- Integration hooks: Tệp tích hợp MXFP4 cho phép sử dụng MXFP4 liền mạch trong framework transformers.
Để kiểm tra xem một mô hình có hỗ trợ MXFP4 hay không, hãy kiểm tra cấu hình của nó:
py from transformers import GptOssConfig
model_id = “openai/gpt-oss-120b” cfg = GptOssConfig.from_pretrained(model_id) print(cfg.quantization_config)
Ví dụ đầu ra:
{
‘modules_to_not_convert’: [
‘model.layers.*.self_attn’,
‘model.layers.*.mlp.router’,
‘model.embed_tokens’,
’lm_head’
],
‘quant_method’: ‘mxfp4’
}
Nếu có 'quant_method': 'mxfp4', mô hình sẽ tự động sử dụng đường dẫn MXFP4 với kernel Triton khi được hỗ trợ.
Nhờ pull request này, bạn có thể tinh chỉnh các mô hình gpt-oss và lưu chúng trực tiếp vào Hub ở định dạng MXFP4, hợp lý hóa việc triển khai với hiệu suất được tối ưu hóa.
Yêu cầu và dự phòng
Để chạy MXFP4 trên GPU, bạn cần:
accelerate,kernelsvàtriton>=3.4đã được cài đặt. Lưu ý rằngPytorch 2.8đã đi kèm vớitriton 3.4, vì vậy bạn chỉ cần cài đặt thủ công triton nếu sử dụngPytorch 2.7.- GPU NVIDIA với compute capability
≥ 7.5. Điều này quay trở lại Tesla, vì vậy bạn có thể chạygpt-oss-20btrên các tầng miễn phí của Google Colab và Kaggle, và trên nhiều GPU tiêu dùng.
Nếu các ràng buộc này không được đáp ứng, transformers sẽ quay trở lại đường dẫn có độ chính xác cao hơn (bfloat16 được sử dụng theo mặc định), yêu cầu khoảng 4 lần bộ nhớ của MXFP4.
Snippet tải GPT-OSS hai lần trên CUDA: một lần với Mxfp4Config(dequantize=True) (tốn nhiều bộ nhớ) và một lần trong đường dẫn lượng tử hóa mặc định (hiệu quả bộ nhớ). Hình 3 hiển thị lượng VRAM được sử dụng sau mỗi lần tải để bạn có thể hình dung các khoản tiết kiệm.
Kernel cho MXFP4
MXFP4 hiệu quả yêu cầu các kernel hiểu các khối 32 phần tử và scale của chúng trong quá trình GEMM và fused ops. Đây là nơi Kernel từ Hub xuất hiện trở lại. transformers tự động kéo vào kernel Triton nhận biết MXFP4 từ kho lưu trữ cộng đồng khi bạn tải một mô hình cần chúng. Kho lưu trữ sẽ xuất hiện trong bộ nhớ cache cục bộ của bạn và sẽ được sử dụng trong quá trình chuyển tiếp. Đối với các kernel MXFP4, người ta không cần sử dụng tham số use_kernels=True như trước đây, nó được đặt thành mặc định trong transformers.
Kiểm tra nhanh chóng với CLI bộ nhớ cache Hugging Face, sau khi chạy gpt-oss-20b trên một GPU tương thích với kernel MXFP4 triton:
shell hf cache scan
Đầu ra mẫu:
shell REPO ID REPO TYPE SIZE ON DISK
kernels-community/triton_kernels model 536.2K openai/gpt-oss-20b model 13.8G
Điều này chỉ ra rằng các kernel MXFP4 đã được tìm nạp và có sẵn để thực thi.
Hãy chạy một số benchmark và xem kernel MXFP4 hoạt động tốt như thế nào. Trong Hình 4, chúng ta thấy rằng các kernel MXFP4 thậm chí còn tốt hơn các kernel MoE và RMSNorm tùy chỉnh cho các batch lớn hơn.
Bạn có thể khám phá và chơi với script benchmark tại đây
Song song hóa Tensor
Song song hóa Tensor (TP) chia các tensor bên trong một lớp trên nhiều GPU (như được hiển thị trong Hình 5). Mỗi GPU nhân shard của nó song song, và sau đó các kết quả một phần được thu thập bằng cách sử dụng các hoạt động all-gather hoặc all-reduce. Điều này làm giảm bộ nhớ trên mỗi GPU và giữ cho tất cả các GPU hoạt động trên cùng một lớp, điều này cải thiện thông lượng khi độ dài chuỗi hoặc batch size tăng lên. TP tốn nhiều giao tiếp và thường hoạt động tốt nhất trên một máy duy nhất với các liên kết intra-node nhanh.
Điều này cho phép điều gì trong transformers
transformers triển khai TP trực tiếp trong from_pretrained. Bạn có thể bắt đầu với plan được xác định trước:
python
chạy với: torchrun –nproc-per-node 4 tp_gpt_oss.py
import torch from transformers import PreTrainedTokenizerFast, GptOssForCausalLM from transformers.distributed import DistributedConfig
model_id = “openai/gpt-oss-120b” tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id) model = GptOssForCausalLM.from_pretrained( model_id, tp_plan=“auto”, # built in TP support dtype=“auto”, ).eval()
messages = [ {“role”: “system”, “content”: “Be concise.”}, {“role”: “user”, “content”: “Explain KV caching briefly.”}, ] inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors=“pt”, return_dict=True, reasoning_effort=“low”, ).to(model.device)
with torch.inference_mode(): generations = model.generate(**inputs, max_new_tokens=128)
print(tokenizer.decode(generations[0][inputs[“input_ids”].shape[-1]:]))
Nếu bạn không có cơ sở hạ tầng để chạy đoạn mã trên, bạn chỉ cần tạo một process trên GPU của chúng tôi bằng cách sử dụng Hugging Face Jobs!
bash
hf jobs run –detach –flavor l4x4 ghcr.io/astral-sh/uv:debian /bin/bash -c
“uv venv .venv –python 3.12 &&
source .venv/bin/activate &&
uv pip install –upgrade torch numpy transformers accelerate triton kernels &&
wget https://huggingface.co/datasets/ariG23498/distributed/raw/main/tp_gpt_oss.py &&
torchrun –nproc-per-node=4 tp_gpt_oss.py”
hf jobscó sẵn cho tất cả người dùng Hugging Face PRO & Enterprise.
Bên dưới, tp_plan="auto" chọn một công thức sharding được xác định trước cho mỗi lớp và kết nối các tập hợp cần thiết. Bạn có thể kiểm tra plan đang hoạt động bằng print(model._tp_plan) nếu bạn muốn xác minh những gì đang được sharded.
Khi nào nên sử dụng TP
Sử dụng TP khi mô hình quá lớn đối với một GPU và bạn muốn tính toán song song, không chỉ vị trí bộ nhớ. TP có xu hướng scale thông lượng với nhiều GPU hơn, đặc biệt đối với các chuỗi dài hoặc các batch lớn hơn.
Nếu bạn tò mò về cách TP khác với
device_map="auto"(vị trí bộ nhớ), câu trả lời Stack Overflow ngắn gọn này giải thích sự khác biệt và khi nào nên sử dụng mỗi loại.
Để tìm hiểu thêm về TP, đây là hai tài nguyên phải đọc:
- Hướng dẫn
transformers: Song song hóa Tensor, các mô hình được hỗ trợ, plan và các điểm mở rộng. - Sổ tay Ultra-Scale: background về TP và mối quan hệ của nó với các chế độ song song hóa khác.
Song song hóa Expert
Expert Parallelism (EP) shard các expert bên trong các lớp MoE trên các GPU. Mỗi token được định tuyến đến một hoặc một vài expert, vì vậy chỉ các expert đó chạy feed-forward pass của chúng. Vì các expert là các MLP độc lập, chúng ta có thể đặt các expert khác nhau trên các rank khác nhau và chỉ trao đổi các trạng thái ẩn cho các token được định tuyến. Điều này giữ cho các phép nhân ma trận còn nguyên vẹn trên mỗi rank và thay thế slicing tensor bằng định tuyến và tập hợp.
Chạy với nhiều process bằng torchrun. EP được bật thông qua cấu hình phân tán và hoạt động với các lớp GPT-OSS MoE ngay lập tức trong transformers.
python
chạy với: torchrun –nproc-per-node 4 ep_gpt_oss.py
import torch from transformers import PreTrainedTokenizerFast, GptOssForCausalLM from transformers.distributed import DistributedConfig
model_id = “openai/gpt-oss-120b” tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id) model = GptOssForCausalLM.from_pretrained( model_id, distributed_config=DistributedConfig(enable_expert_parallel=True), # enabling EP dtype=“auto”, ).eval()
messages = [ {“role”: “system”, “content”: “Be concise.”}, {“role”: “user”, “content”: “Explain KV caching briefly.”}, ] inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors=“pt”, return_dict=True, reasoning_effort=“low”, ).to(model.device)
with torch.inference_mode(): generations = model.generate(**inputs, max_new_tokens=128)
print(tokenizer.decode(generations[0][inputs[“input_ids”].shape[-1]:]))
Đây là cách bạn sẽ chạy bằng hf jobs
bash
hf jobs run –detach –flavor l4x4 ghcr.io/astral-sh/uv:debian /bin/bash -c
“uv venv .venv –python 3.12 &&
source .venv/bin/activate &&
uv pip install –upgrade torch numpy transformers accelerate triton kernels &&
wget https://huggingface.co/datasets/ariG23498/distributed/raw/main/ep_gpt_oss.py &&
torchrun –nproc-per-node=4 ep_gpt_oss.py”
Khi bạn bật Expert Parallelism, Tensor Parallelism cũng được kích hoạt. Điều này có nghĩa là bạn tận hưởng những điều tốt nhất của cả hai thế giới!
Lớp Cửa sổ Trượt Động & Bộ nhớ Cache
Nhiều LLM gần đây sử dụng attention cửa sổ trượt, hoặc sự kết hợp của các lớp attention trượt và toàn cục, như một phương tiện để tiết kiệm bộ nhớ và giảm các matmul bậc hai tốn kém phát triển theo độ dài chuỗi. Tuy nhiên, việc triển khai bộ nhớ cache KV động trong transformers vẫn tiếp tục phân bổ không gian theo độ dài chuỗi, mà không xem xét các lớp attention riêng lẻ. Bạn luôn có thể tối ưu hóa bộ nhớ bằng cách sử dụng biên dịch (nghĩa là, các hình dạng cố định), nhưng đó là một kịch bản riêng biệt.
transformers hiện có một DynamicSlidingWindowLayer và một DynamicCache nhận biết cấu hình. Nếu cấu hình mô hình khai báo attention cửa sổ trượt hoặc attention kết hợp (cả các lớp attention trượt và toàn cục đều được sử dụng), bộ nhớ cache sẽ ngừng phát triển vượt quá cửa sổ cho các lớp trượt. Nếu bạn không chuyển cấu hình, hành vi vẫn như trước (KV đầy đủ, luôn phát triển khi độ dài chuỗi tăng lên).
Đối với các mô hình chỉ sử dụng các lớp cửa sổ trượt, chẳng hạn như Mistral 7B, bộ nhớ cache ngừng phát triển khi chuỗi đạt đến kích thước cửa sổ (4096, trong trường hợp này). Điều này có ý nghĩa, vì các lớp trượt không thể nhìn xa hơn 4K token trước đó.
OpenAI gpt-oss luân phiên giữa các lớp attention trượt và toàn cục, dẫn đến tổng bộ nhớ cache KV giảm một nửa, như chúng ta sẽ thấy, khi độ dài chuỗi tăng lên. Điều này cung cấp cho chúng ta:
- Bộ nhớ KV-cache thấp hơn nhiều cho các mô hình có attention trượt hoặc kết hợp (ví dụ: GPT-OSS). Tốc độ tăng trưởng bộ nhớ cache ổn định sau khi đạt đến cửa sổ (ví dụ: 4K cho Mistral; 128 cho các lớp trượt GPT-OSS), thay vì scale tuyến tính với tổng số token được tạo. (GitHub, Transformers)
- Chiến thắng về tốc độ/độ trễ trên các prompt dài/thế hệ dài: các tensor KV nhỏ hơn có nghĩa là đọc/ghi attention nhẹ hơn và ít áp lực băng thông bộ nhớ hơn, đặc biệt là sau khi đạt đến cửa sổ. (Đây là động lực trung tâm đằng sau LLM cửa sổ trượt/kết hợp.) (AI21, Blog vLLM)
Cách sử dụng nó
Bộ nhớ cache được tối ưu hóa được đặt theo mặc định, điều đó có nghĩa là bạn không phải thực hiện bất kỳ thay đổi nào đối với mã hiện có của mình. Nếu bạn muốn tạo DynamicCache một cách rõ ràng, đây là cách bạn sẽ làm:
python from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
model_id = “openai/gpt-oss-20b”
tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained( model_id, dtype=“auto”, device_map=“auto”, ).eval()
messages = [ {“role”: “system”, “content”: “Always respond in riddles”}, {“role”: “user”, “content”: “What is the weather like in Madrid?”}, ]
inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors=“pt”, return_dict=True, reasoning_effort=“low”, ).to(model.device)
cache = DynamicCache(config=model.config) # tạo bộ nhớ cache với cấu hình của mô hình
generated = model.generate( **inputs, max_new_tokens=500, past_key_values=cache ) print(tokenizer.decode(generated[0][inputs[“input_ids”].shape[-1]:]))
Hình 6 cho thấy sự khác biệt lớn như thế nào đối với chúng ta khi sử dụng Dynamic KV Cache với attention cửa sổ trượt.
Batching liên tục & Chú ý theo Trang
Một quá trình tạo tự hồi quy điển hình trông giống như Hình 7. Bạn nhập các token prefill và mô hình dự đoán từng token mới một sau một cho đến khi nó dự đoán token EOS (End of Sequence).
Hãy xem quá trình tạo trông như thế nào khi chúng ta chuyển một batch các đầu vào. Trong Hình 8, bạn nhận thấy rằng một số thế hệ kết thúc sớm hơn những thế hệ khác. Sự không phù