TransformerEngine
v1.13
Quickstart | การติดตั้ง คู่มือผู้ใช้ | ตัวอย่าง FP8 Convergence | การบูรณาการ บันทึกย่อ
Transformer Engine (TE) เป็นห้องสมุดสำหรับเร่งรุ่นหม้อแปลงใน NVIDIA GPU รวมถึงการใช้จุดลอยตัว 8 บิต (FP8) ความแม่นยำบน GPU Hopper เพื่อให้ประสิทธิภาพที่ดีขึ้น TE จัดเตรียมชุดของการสร้างบล็อกที่มีประสิทธิภาพสูงสำหรับสถาปัตยกรรมหม้อแปลงที่ได้รับความนิยมและ API ที่มีความแม่นยำแบบผสมอัตโนมัติซึ่งสามารถใช้งานได้อย่างราบรื่นด้วยรหัสเฉพาะเฟรมของคุณ TE ยังรวมถึง Framework Agnostic C ++ API ที่สามารถรวมเข้ากับห้องสมุดการเรียนรู้ลึกอื่น ๆ เพื่อเปิดใช้งานการสนับสนุน FP8 สำหรับหม้อแปลง
เนื่องจากจำนวนพารามิเตอร์ในโมเดลหม้อแปลงยังคงเติบโตการฝึกอบรมและการอนุมานสำหรับสถาปัตยกรรมเช่น Bert, GPT และ T5 กลายเป็นหน่วยความจำและการคำนวณมาก กรอบการเรียนรู้ที่ลึกที่สุดฝึกอบรมด้วย FP32 โดยค่าเริ่มต้น อย่างไรก็ตามสิ่งนี้ไม่จำเป็นเพื่อให้เกิดความแม่นยำอย่างเต็มที่สำหรับรูปแบบการเรียนรู้ที่ลึกซึ้งมากมาย การใช้การฝึกอบรมแบบผสมผสานซึ่งรวมความแม่นยำเดี่ยว (FP32) กับรูปแบบความแม่นยำต่ำกว่า (เช่น FP16) เมื่อฝึกอบรมแบบจำลองส่งผลให้เกิดการเร่งความเร็วอย่างมีนัยสำคัญด้วยความแม่นยำน้อยที่สุดเมื่อเทียบกับการฝึกอบรม FP32 ด้วยสถาปัตยกรรม Hopper GPU FP8 Precision ได้รับการแนะนำซึ่งให้ประสิทธิภาพที่ดีขึ้นกว่า FP16 โดยไม่มีการย่อยสลายความแม่นยำ แม้ว่ากรอบการเรียนรู้เชิงลึกที่สำคัญทั้งหมดสนับสนุน FP16 แต่การสนับสนุน FP8 นั้นไม่สามารถใช้งานได้ในกรอบงานในปัจจุบัน
TE จัดการปัญหาการสนับสนุน FP8 โดยการจัดหา APIs ที่รวมเข้ากับห้องสมุดรูปแบบภาษาขนาดใหญ่ (LLM) ยอดนิยม (LLM) มันมี Python API ซึ่งประกอบด้วยโมดูลเพื่อสร้างเลเยอร์หม้อแปลงได้อย่างง่ายดายรวมถึงห้องสมุดเฟรมเวิร์ก-เครื่องผืดใน C ++ รวมถึงโครงสร้างและเคอร์เนลที่จำเป็นสำหรับการสนับสนุน FP8 โมดูลที่จัดทำโดย TE ภายในรักษาปัจจัยการปรับสเกลและค่าอื่น ๆ ที่จำเป็นสำหรับการฝึกอบรม FP8 ทำให้การฝึกอบรมแบบผสมผสานแบบผสมง่ายขึ้นอย่างมากสำหรับผู้ใช้
import torch
import transformer_engine . pytorch as te
from transformer_engine . common import recipe
# Set dimensions.
in_features = 768
out_features = 3072
hidden_size = 2048
# Initialize model and inputs.
model = te . Linear ( in_features , out_features , bias = True )
inp = torch . randn ( hidden_size , in_features , device = "cuda" )
# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe . DelayedScaling ( margin = 0 , fp8_format = recipe . Format . E4M3 )
# Enable autocasting for the forward pass
with te . fp8_autocast ( enabled = True , fp8_recipe = fp8_recipe ):
out = model ( inp )
loss = out . sum ()
loss . backward () import flax
import jax
import jax . numpy as jnp
import transformer_engine . jax as te
import transformer_engine . jax . flax as te_flax
from transformer_engine . common import recipe
BATCH = 32
SEQLEN = 128
HIDDEN = 1024
# Initialize RNG and inputs.
rng = jax . random . PRNGKey ( 0 )
init_rng , data_rng = jax . random . split ( rng )
inp = jax . random . normal ( data_rng , [ BATCH , SEQLEN , HIDDEN ], jnp . float32 )
# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe . DelayedScaling ( margin = 0 , fp8_format = recipe . Format . HYBRID )
# Enable autocasting for the forward pass
with te . fp8_autocast ( enabled = True , fp8_recipe = fp8_recipe ):
model = te_flax . DenseGeneral ( features = HIDDEN )
def loss_fn ( params , other_vars , inp ):
out = model . apply ({ 'params' : params , ** other_vars }, inp )
return jnp . mean ( out )
# Initialize models.
variables = model . init ( init_rng , inp )
other_variables , params = flax . core . pop ( variables , 'params' )
# Construct the forward and backward function
fwd_bwd_fn = jax . value_and_grad ( loss_fn , argnums = ( 0 , 1 ))
for _ in range ( 10 ):
loss , ( param_grads , other_grads ) = fwd_bwd_fn ( params , other_variables , inp )วิธีที่เร็วที่สุดในการเริ่มต้นด้วยเครื่องยนต์หม้อแปลงคือการใช้ภาพ Docker ในแคตตาล็อก Nvidia GPU Cloud (NGC) ตัวอย่างเช่นการใช้คอนเทนเนอร์ NGC Pytorch แบบโต้ตอบ
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:23.10-py3โดยที่ 23.10 เป็นเวอร์ชันคอนเทนเนอร์ ตัวอย่างเช่น 23.10 สำหรับการเปิดตัวตุลาคม 2566
ในการติดตั้ง Transformer Engine เวอร์ชันที่เสถียรล่าสุด
pip install git+https://github.com/NVIDIA/TransformerEngine.git@stableสิ่งนี้จะตรวจจับโดยอัตโนมัติหากมีการติดตั้งเฟรมเวิร์กการเรียนรู้ลึกที่รองรับและสร้างการสนับสนุนเครื่องยนต์หม้อแปลงสำหรับพวกเขา หากต้องการระบุเฟรมเวิร์กอย่างชัดเจนให้ตั้งค่าตัวแปรสภาพแวดล้อม NVTE_FRAMEWORK เป็นรายการที่คั่นด้วยเครื่องหมายจุลภาค (เช่น NVTE_FRAMEWORK = JAX, PYTORCH, พาย)
อีกทางเลือกหนึ่งแพ็คเกจสามารถติดตั้งโดยตรงจาก PYPI ของ Transformer Engine เช่น
pip install transformer_engine[pytorch]เพื่อให้ได้การผูก python ที่จำเป็นสำหรับเครื่องยนต์หม้อแปลงต้องระบุเฟรมเวิร์กที่จำเป็นอย่างชัดเจนว่าเป็นการพึ่งพาพิเศษในรายการที่คั่นด้วยเครื่องหมายจุลภาค (เช่น [Jax, Pytorch, Paddle]) เครื่องยนต์หม้อแปลงเรือล้อสำหรับห้องสมุดหลักเช่นเดียวกับส่วนขยาย Paddlepaddle การแจกแจงแหล่งที่มาถูกจัดส่งสำหรับส่วนขยาย JAX และ Pytorch
ดูคู่มือการติดตั้ง
Transformer Engine Release V0.11.0 เพิ่มการสนับสนุนสำหรับ Flashattention-2 ใน Pytorch เพื่อการปรับปรุงประสิทธิภาพ
เป็นปัญหาที่ทราบกันดีว่าการรวบรวม Flashattention-2 นั้นใช้ทรัพยากรมากและต้องใช้ RAM จำนวนมาก (ดูข้อผิดพลาด) ซึ่งอาจนำไปสู่ข้อผิดพลาดของหน่วยความจำในระหว่างการติดตั้งเครื่องยนต์หม้อแปลง โปรดลองตั้งค่า max_jobs = 1 ในสภาพแวดล้อมเพื่อหลีกเลี่ยงปัญหา
โปรดทราบว่า NGC Pytorch 23.08+ คอนเทนเนอร์รวมถึง Flashattention-2
ในความพยายามที่จะรวมคำจำกัดความและการใช้หน้ากากความสนใจในทั้งสามเฟรมเวิร์กในเครื่องยนต์หม้อแปลงหน้ากากพลาดิ้งได้เปลี่ยนจากความหมายที่แท้จริงรวมถึงตำแหน่งที่สอดคล้องกันเพื่อให้ความสนใจกับการยกเว้นตำแหน่งนั้นในการใช้งาน pytorch ของเรา เนื่องจาก v1.7 ประเภทหน้ากากความสนใจทั้งหมดเป็นไปตามคำจำกัดความเดียวกันซึ่งหมายถึงการปิดบังตำแหน่งที่สอดคล้องกันและวิธีการที่ผิดพลาดรวมถึงตำแหน่งนั้นในการคำนวณความสนใจ
ตัวอย่างของการเปลี่ยนแปลงนี้คือ
# for a batch of 3 sequences where `a`s, `b`s and `c`s are the useful tokens
# and `0`s are the padding tokens,
[a, a, a, 0, 0,
b, b, 0, 0, 0,
c, c, c, c, 0]
# the padding mask for this batch before v1.7 is,
[ True, True, True, False, False,
True, True, False, False, False,
True, True, True, True, False]
# and for v1.7 onwards it should be,
[False, False, False, True, True,
False, False, True, True, True,
False, False, False, False, True]FP8 ได้รับการทดสอบอย่างกว้างขวางในสถาปัตยกรรมแบบจำลองและการกำหนดค่าที่แตกต่างกันและเรา ไม่พบความแตกต่างอย่างมีนัยสำคัญ ระหว่างเส้นโค้งการสูญเสียการฝึกอบรม FP8 และ BF16 FP8 ยังได้รับการตรวจสอบความถูกต้องในงาน LLM ดาวน์สตรีม (เช่น Lambada และ Wikitext) ด้านล่างนี้เป็นตัวอย่างของแบบจำลองที่ทดสอบสำหรับการบรรจบกันในกรอบที่แตกต่างกัน
| แบบอย่าง | กรอบ | แหล่งที่มา |
|---|---|---|
| T5-770M | jax/t5x | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/t5x#convergence-and-performance |
| mpt-1.3b | นักแต่งเพลงโมเสค | https://www.mosaicml.com/blog/coreweave-nvidia-h100-part-1 |
| GPT-5B | jax/paxml | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/pax#h100-results |
| GPT-5B | เฟรมเวิร์ก Nemo | มีตามคำขอ |
| LLAMA2-7B | อาลีบาบาปาย | https://mp.weixin.qq.com/s/nqt0ukxlbxyh5031zbdebq |
| T5-11B | jax/t5x | มีตามคำขอ |
| mpt-13b | นักแต่งเพลงโมเสค | https://www.databricks.com/blog/turboarged-training-optimizing-databricks-mosaic-ai-stack-fp8 |
| GPT-22B | เฟรมเวิร์ก Nemo | มีตามคำขอ |
| LLAMA2-70B | อาลีบาบาปาย | https://mp.weixin.qq.com/s/nqt0ukxlbxyh5031zbdebq |
| GPT-175B | jax/paxml | https://github.com/nvidia/jax-toolbox/tree/main/rosetta/rosetta/projects/pax#h100-results |
เครื่องยนต์หม้อแปลงได้รับการรวมเข้ากับเฟรมเวิร์ก LLM ยอดนิยมเช่น:
เรายินดีต้อนรับการมีส่วนร่วมกับ Transformer Engine! หากต้องการมีส่วนร่วมในเครื่องยนต์หม้อแปลงและทำการร้องขอการดึงให้ทำตามแนวทางที่ระบุไว้ในคู่มือการสนับสนุน
