Hướng dẫn Triton RMSNorm: kernel, backward và…

Có thể bạn từng tối ưu một xử lý web phải duyệt cùng một bảng nhiều lần. Ý tưởng khá đơn giản: gom các thao tác để tránh công việc trung gian. Nhưng trước khi so sánh thời gian, cần kiểm tra phiên bản mới vẫn cho cùng kết quả.

Chúng ta sẽ áp dụng cách làm này bên trong một mạng nơ-ron. Phép toán có tên RMSNorm: nó đưa các dãy số về cùng thang đo. Ta bắt đầu bằng phiên bản PyTorch dễ đọc, rồi gộp phép tính vào một kernel Triton, hàm chạy trên GPU. Khi huấn luyện, đầu ra đúng vẫn chưa đủ: cần tính đúng cả gradient.

Tải toàn bộ dự án

Lộ trình tạo ra RMSNorm kết nối với Autograd, được dùng trong một mạng nhỏ thực sự được huấn luyện rồi nạp lại. Mục tiêu là kiểm tra phép tính trong một phạm vi xác định, không phải chứng minh “Triton luôn nhanh hơn”. Ta sẽ xem riêng độ đúng số học và độ trễ.

Vấn đề và công thức tham chiếu

Chính xác thì ta chuẩn hóa gì?

Một mạng biến đầu vào thành các dãy số trung gian: đó là activation. Chúng thay đổi theo từng ví dụ. Trọng số là các tham số mà quá trình huấn luyện điều chỉnh để thay đổi hành vi mạng.

Ở đây, chuẩn hóa activation là kiểm soát thang đo của chúng. Ta không sửa chuỗi ký tự hay định dạng tài liệu: ta áp dụng một biến đổi số học bên trong mô hình.

Trong dự án này, đầu vào x là bảng hai chiều, gọi là tensor. Tensor có thể có nhiều chiều hơn, nhưng ở đây ta dùng ma trận M hàng và N cột. Mỗi hàng được chuẩn hóa độc lập.

Đối tượng	Dạng trong hướng dẫn	Cách hiểu
`x`	`M × N`	`M` hàng activation, mỗi hàng có `N` số.
`gamma`	`N`	Một trọng số mỗi cột, dùng chung cho mọi hàng.
`r`	Một số mỗi hàng	Nghịch đảo của thước đo thang đo tính cho hàng đó.
`y`	`M × N`	Activation sau chuẩn hóa và nhân với `gamma`.

RMS là root mean square, tức “căn bậc hai của trung bình bình phương”. Bình phương tránh cho giá trị dương và âm triệt tiêu nhau. Trung bình tính đến độ dài hàng; căn đưa ta về một thước đo thang đo.

Với một hàng, công thức tham chiếu là:

r = 1 / sqrt(mean(x²) + epsilon)

y = x × r × gamma

Các phép nhân thực hiện theo từng tọa độ. r chung cho cả hàng, còn mỗi cột có gamma riêng. Số nhỏ epsilon, đặt là 1e-6 trong các hàm dưới đây, được thêm trước căn bậc hai. Nó đặc biệt tránh chia cho 0 với hàng toàn số 0.

Tính với hai giá trị

Lấy x = [3, 4], gamma = [1, 1] và epsilon = 1e-6. Đây là minh họa số học, không phải một hàng lấy từ phép đo GPU.

Bước	Tính toán	Kết quả xấp xỉ
Bình phương	`[3², 4²]`	`[9, 16]`
Trung bình bình phương	`(9 + 16) / 2`	`12,5`
Mẫu số ổn định	`sqrt(12,5 + 0,000001)`	`3,535534`
Nghịch đảo `r`	`1 / 3,535534`	`0,282843`
Đầu ra	`[3 × r, 4 × r]`	`[0,848528 ; 1,131371]`

Ta giữ tỷ lệ giữa hai giá trị nhưng thay đổi thang đo. Đầu ra thứ hai lớn hơn 1: chuẩn hóa không có nghĩa là giới hạn từng số trong khoảng 0 đến 1.

Chỉ đổi gamma thành [2, 0.5]. Đầu ra trở thành khoảng [1,697056 ; 0,565685]. Các trọng số này giúp mạng điều chỉnh riêng tầm quan trọng của từng tọa độ sau khi đổi thang đo; chúng không chỉ là trang trí.

Khác với LayerNorm, RMSNorm không trừ trung bình của hàng. Trong ví dụ này, hai đầu ra vẫn dương nên trung bình của chúng không bằng 0. Bài báo RMSNorm gốc nêu sự khác biệt này; tài liệu RMSNorm của PyTorch cũng đưa công thức. Tuy vậy, so sánh của ta dùng tham chiếu viết trong hướng dẫn, với lựa chọn độ chính xác và epsilon rõ ràng.

Vì sao cần một lượt đi và một lượt về?

Forward, hay lượt đi, tính y từ x và gamma. Khi huấn luyện, mạng tạo ra dự đoán, được so với mục tiêu bằng một loss: con số cần giảm.

Backward, hay lan truyền ngược, truyền gradient theo chiều ngược lại. Gradient đo độ nhạy của loss trước một thay đổi nhỏ của giá trị. Nó không tự cập nhật trọng số: optimizer dùng nó để làm việc đó.

Ở đầu vào backward của RMSNorm, dy là “gradient của loss theo y”. Nó không phải hiệu y - x, cũng không trực tiếp là sai số giữa dự đoán và mục tiêu. Các phép toán sau RMSNorm đã góp phần vào phép tính này.

Ta phải trả về hai đối tượng: dx để tiếp tục lan truyền ngược tới lớp trước, và dgamma để học trọng số chuẩn hóa. Với u = dy × gamma, các biểu thức cần tìm là:

dx = r × u − x × r³ × mean(u × x)

dgamma = sum_lignes(dy × x × r)

Ta sẽ giải thích các công thức khi triển khai. Hãy nhớ khác biệt: trung bình trong dx lấy trên N cột của một hàng, còn tổng của dgamma lấy trên M hàng.

PyTorch cung cấp thao tác tensor và Autograd tổ chức đạo hàm của chúng. Triton cho phép viết tính toán GPU với cú pháp gần Python. Kernel gom các thao tác chạy trên card; nó không phải mô hình huấn luyện sẵn. Ở đây ta phải cung cấp rõ backward để nối nó với phần còn lại của mạng.

Điều kiện Docker và GPU

Bạn cần biết đọc một hàm Python, thao tác bảng và dùng terminal với Docker. Không cần kinh nghiệm lập trình CUDA trước đó để theo dõi giải thích.

Lộ trình đã xác thực yêu cầu Linux x86_64, Docker với Compose, NVIDIA Container Toolkit và RTX 3090 24 Gio. Cần 16 Gio RAM và 16 Gio đĩa, không tính lớp Docker dùng chung. Bộ nhớ video GPU (VRAM) khác RAM hệ thống. Không có cấu hình CPU nào thay thế được việc kernel chạy CUDA.

Giải nén và đặt terminal ở thư mục gốc. Mọi lệnh sau đều chạy từ cùng thư mục này. Trích đoạn chỉ hiển thị hàm cốt lõi, không phải tệp độc lập: import, hợp đồng đầy đủ và hàm phụ nằm trong dự án tải về.

setup.sh tạo thư mục cục bộ và chỉ ghi UID/GID vào .env: các định danh này quản lý quyền tệp. acquire ở đây không tải mô hình huấn luyện sẵn hay dữ liệu ngoài. Dữ liệu là tổng hợp; outputs/ chứa kiểm tra và artifacts/ chứa checkpoint, tức trạng thái mạng đã lưu. Phụ thuộc được giải quyết và kiểm tra khi build.

sh scripts/setup.sh
docker compose config --quiet
docker compose build acquire gpu
docker compose run --rm acquire python scripts/acquire.py --config configs/minimum.yaml
docker compose run --rm gpu python scripts/check_environment.py --config configs/minimum.yaml

docker compose config --quiet kiểm tra cấu hình Compose. Build chuẩn bị dịch vụ acquire và gpu. Cuối cùng, check_environment.py khởi tạo CUDA trong container và thực hiện một phép tính: thấy card từ Linux chưa đủ chứng minh Docker dùng được nó.

Nếu kiểm tra cuối thất bại, dừng ở bước này và kiểm tra tích hợp theo hướng dẫn NVIDIA Container Toolkit. Đừng đổi công thức RMSNorm để né vấn đề môi trường. Tái sử dụng lớp Docker có thể đổi thời gian build; nó không đổi phép tính số học cần so sánh.

Kiểm tra các hàng rào trước kernel

Trước khi tối ưu một hàm web, bạn sẽ giữ các test đầu vào và đầu ra. Ta làm tương tự với tính toán GPU, đồng thời kiểm tra gradient và tổ chức bộ nhớ.

Bộ test công khai bao phủ kích thước bất thường, hợp đồng bộ nhớ, trường hợp âm và kiểm tra gradient. Trường hợp âm là đầu vào phải bị từ chối hoặc biến thể sai mà test phải phát hiện. Một lỗi dự kiến được phát hiện đúng là test thành công.

docker compose run --rm gpu python -m pytest -q

Bộ test phải kết thúc không lỗi trước khi đo hay huấn luyện. Thành công này áp dụng cho mã và các trường hợp của dự án; nó không biến toán tử thành kernel phổ quát chấp nhận mọi kích thước, kiểu số hay bố trí bộ nhớ.

Tham chiếu PyTorch: oracle số học

Oracle kiểm thử đơn giản là phiên bản dùng để so sánh với triển khai đang xét. Tham chiếu của ta ưu tiên dễ đọc: có thể đối chiếu từng dòng Python với công thức trên.

dtype là kiểu số của tensor. FP16, FP32 và FP64 là số thực dấu phẩy động lưu trên 16, 32 và 64 bit; độ chính xác khác nhau. Tham chiếu dùng float64 nếu đầu vào đã là float64, nếu không dùng FP32, rồi chuyển đầu ra về kiểu đầu vào.

def rmsnorm_reference(x, gamma, epsilon=1e-6):
    dtype = torch.float64 if x.dtype == torch.float64 else torch.float32
    xf = x.to(dtype)
    gf = gamma.to(dtype)
    r = torch.rsqrt(xf.square().mean(-1, keepdim=True) + epsilon)
    return (xf * r * gf).to(x.dtype)


def manual_backward(x, gamma, dy, epsilon=1e-6):
    r = torch.rsqrt(x.square().mean(-1, keepdim=True) + epsilon)
    u = dy * gamma
    dx = r * u - x * r.pow(3) * (u * x).mean(-1, keepdim=True)
    dgamma = (dy * x * r).sum(0)
    return dx, dgamma

Đọc forward từng dòng

xf và gf là phiên bản làm việc của x và gamma, đổi sang kiểu được chọn. square() tính bình phương. mean(-1, keepdim=True) lấy trung bình trên trục cuối, tức các cột.

Vì sao giữ chiều bằng keepdim=True? Với đầu vào M × N, trung bình có dạng M × 1. Mỗi hàng giữ số riêng để PyTorch nhân với mọi cột cùng hàng. Việc tự mở rộng hình dạng này gọi là broadcasting. gamma có dạng N được dùng lại trên mọi hàng.

torch.rsqrt(a) tính 1 / sqrt(a). Vì vậy dòng cuối áp dụng đúng x × r × gamma, rồi trả về dtype ban đầu. Tính tạm bằng FP32 không có nghĩa đầu ra sẽ còn FP32 nếu đầu vào là FP16.

Công thức backward độc lập, dễ đọc

manual_backward() biểu diễn hai đạo hàm mà không qua Triton. Mục đích là đối chiếu công thức tường minh với gradient do Autograd tạo ra.

Khác với rmsnorm_reference(), hàm này không chuyển rõ đầu vào sang kiểu chính xác hơn. Vì vậy phải đọc nó theo dtype đối số; riêng nó không phải công thức giảm FP32 cho tensor FP16. Kernel sẽ cho thấy chuyển đổi này trong lúc tải.

gradcheck thêm một kiểm tra: nó nhiễu nhẹ đầu vào và so sánh biến thiên số học của đầu ra với đạo hàm tính được. Ý tưởng là độ dốc của hàm vô hướng có thể ước tính bằng (f(a + h) - f(a - h)) / (2 × h), với h nhỏ. Tài liệu gradcheck nêu vai trò của độ chính xác kép trong thiết lập thông thường.

Trong lộ trình này, gradcheck float64 áp dụng cho tham chiếu PyTorch, không phải kernel Triton FP32/FP16. Kernel sẽ được so sánh trong các kiểm tra riêng. Phân biệt này tránh gán cho một chương trình test thực hiện trên chương trình khác.

Bây giờ hãy tạo kiểm tra tham chiếu này và ba tập phân vùng tổng hợp:

docker compose run --rm gpu python -m rmsnorm_kernel.cli reference --config configs/minimum.yaml

Các phân vùng tách dữ liệu cho huấn luyện, xác thực và test. Chi tiết cách tạo thuộc về dự án; không cần dữ liệu nghiệp vụ hay mô hình ngoài để nghiên cứu phép toán này.

Forward Triton: một hàng, một chương trình

Hợp đồng bộ nhớ là một phần của phép tính

Hãy hình dung một hàm nhận bảng chữ nhật nhưng duyệt bộ nhớ như thể mọi hàng xếp liền nhau. Nó sẽ sai với một view có phần tử cách nhau khác, dù kích thước hiển thị có vẻ đúng.

Tensor contiguous ở đây có bố trí gọn mà kernel mong đợi. Stride mô tả bước trong bộ nhớ để đi qua một hàng hay cột. Ví dụ, chuyển vị có thể tạo view không contiguous mà không phải di chuyển mọi số.

Trước khi ghi đầu ra, hợp đồng dự án từ chối đầu vào không contiguous, dtype ngoài FP32/FP16, thiết bị khác nhau, epsilon không hợp lệ, N lớn hơn 4096 và vùng tạm hay scratch lớn hơn 64 Mio. Phần thân check_contract() không được chép lại: hãy giữ các kiểm tra trong tệp đầy đủ.

@triton.jit
def rmsnorm_forward_kernel(
    X, G, Y, R, N: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr
):
    row = tl.program_id(0)
    column = tl.arange(0, BLOCK)
    x = tl.load(X + row * N + column, column < N, other=0).to(tl.float32)
    gamma = tl.load(G + column, column < N, other=0).to(tl.float32)
    inverse = tl.rsqrt(tl.sum(x * x, 0) / N + EPS)
    y = x * inverse * gamma
    tl.store(Y + row * N + column, y, column < N)
    tl.store(R + row, inverse)


def forward_op(x, gamma, epsilon=1e-6):
    check_contract(x, gamma, epsilon)
    m, n = x.shape
    y = torch.empty_like(x)
    inverse = torch.empty(m, device=x.device, dtype=torch.float32)
    rmsnorm_forward_kernel[(m,)](
        x, gamma, y, inverse, n, epsilon, triton.next_power_of_2(n)
    )
    return y, inverse

Ai tính hàng nào?

Phần Python forward_op() kiểm tra đầu vào, cấp phát đầu ra, rồi chạy kernel với lưới [(m,)]. Vì vậy ta yêu cầu M thể hiện chương trình Triton, một cho mỗi hàng.

tl.program_id(0) cho mỗi thể hiện số hàng của nó. Số này không phải tiến trình Docker hay thêm card đồ họa. tl.arange(0, BLOCK) tạo vị trí cột được xử lý cùng lúc. Hướng dẫn Triton đầu tiên trình bày mô hình này và các mặt nạ bộ nhớ.

X + row * N + column chỉ các giá trị của hàng cần lấy. Công thức dùng N vì hàng thật có N giá trị; không dùng kích thước làm tròn BLOCK để sang hàng tiếp theo.

Decorator @triton.jit giao Triton biên dịch kernel. Tham số tl.constexpr được biết khi chuyên biệt hóa: chúng mô tả kích thước block và hằng của phép tính, không phải trọng số đã học.

N và BLOCK: hai kích thước không được nhầm

triton.next_power_of_2(n) chọn độ rộng làm việc bằng lũy thừa hai kế tiếp. Với N=257, chương trình thao tác BLOCK=512 vị trí, nhưng chỉ cột 0 đến 256 thuộc hàng.

Điều đó không biến tensor thành ma trận 512 cột. Vị trí thêm là chỗ tính cần che. tl.load(..., column < N, other=0) không đọc ngoài hàng mà thay bằng 0; tl.store(..., column < N) chỉ ghi giá trị hợp lệ.

Vì sao là 0? Nó không đổi tổng bình phương. Tuy nhiên, không được tính nó như quan sát bổ sung trong trung bình. Tổng phải chia cho N, không phải BLOCK.

Lấy ví dụ nhỏ không phải phép đo: hàng [1, 1, 1] xử lý với block bốn vị trí. Tổng bình phương [1, 1, 1, 0] vẫn là 3. Chia 3 cho trung bình đúng là 1; chia 4 thành 0,75. Mặt nạ bảo vệ truy cập bộ nhớ nhưng không chọn mẫu số đúng thay ta.

Lỗi này không lộ ở N=256 vì khi đó BLOCK cũng bằng 256. Đó là lý do kích thước bất thường là test hữu ích, không chỉ là chi tiết trình bày.

Vì sao đổi kiểu trước khi reduce?

Reduction biến nhiều giá trị thành một kết quả, như tổng hay trung bình. Ở đây số được chuyển FP32 ngay khi tải, trước bình phương và cộng tổng. Chuyển đổi sau một phép tính đã làm tròn không thể lấy lại độ chính xác đã mất.

Kernel tính RMS nghịch đảo, áp dụng gamma, rồi ghi y. Y có cùng kiểu với x nhờ torch.empty_like(x); R giữ một nghịch đảo mỗi hàng, rõ ràng ở FP32, cho backward. empty_like chỉ cấp phát mà không điền 0: kernel phải ghi mọi đầu ra hữu ích.

Ta gộp forward nhưng không loại bỏ mọi lưu trữ trung gian: R được cố ý giữ lại giữa lượt đi và lượt về.

docker compose run --rm gpu python -m rmsnorm_kernel.cli check-forward --config configs/minimum.yaml

Kiểm tra này so sánh forward với các kích thước dự kiến, gồm N không phải lũy thừa hai. Bài LayerNorm của Triton giải thích thêm kernel chuẩn hóa và reduction; nó không thay thế hợp đồng RMSNorm hẹp hơn của dự án.

Backward đầu vào: giữ đúng N

Khó khăn: mỗi giá trị cũng ảnh hưởng mẫu số

Để tính dx, không thể chỉ nhân dy với r × gamma, vì r cũng phụ thuộc toàn bộ hàng x.

Trở lại [3, 4]. Nếu đổi số đầu tiên, tử số đổi nhưng tổng bình phương cũng đổi. Mẫu số chung của hai đầu ra do đó thay đổi. Backward phải tính đến tác động tập thể này.

Đó là vai trò hai hạng của công thức: r × u mô tả tác động trực tiếp; −x × r³ × mean(u × x) thêm tác động của thay đổi chuẩn hóa. u = dy × gamma trước hết đưa trọng số theo cột vào.

Minh họa cùng ví dụ: đặt lại gamma = [1, 1] và chọn gradient đi lên dy = [1, 2]. Ta có u = [1, 2], rồi mean(u × x) = (3 + 8) / 2 = 5,5. Với epsilon = 1e-6, công thức cho xấp xỉ:

dx = [-0,090510 ; 0,067882]

Gradient đầu tiên âm dù dy dương. Đây không phải bất thường: tăng một đầu vào có thể giảm đầu ra khác do tăng mẫu số. Gradient tính mọi đường mà đầu vào này ảnh hưởng loss. Các số này minh họa công thức, không phải chênh lệch đo của kernel.

Lập phương của r từ đâu?

Dẫn xuất nhỏ này không bắt buộc để chạy dự án, nhưng làm mã bớt bí ẩn.

Đặt s = mean(x²) + epsilon. Vì r = s^(-1/2), đạo hàm của r theo s là −(1/2) × s^(-3/2), tức −r³ / 2. Đạo hàm của s theo tọa độ x_i là 2 × x_i / N.

Nhân hai biểu thức ta được ∂r/∂x_i = −x_i × r³ / N. Còn phải cộng đóng góp của mọi đầu ra dùng cùng r, tạo nên hạng thứ hai của dx.

Vì vậy, lập phương không phải tham số điều chỉnh đến khi test qua: nó đến từ đạo hàm của nghịch đảo căn bậc hai.

@triton.jit
def rmsnorm_dx_kernel(X, G, DY, R, DX, PARTIAL, N: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    column = tl.arange(0, BLOCK)
    x = tl.load(X + row * N + column, column < N, other=0).to(tl.float32)
    gamma = tl.load(G + column, column < N, other=0).to(tl.float32)
    dy = tl.load(DY + row * N + column, column < N, other=0).to(tl.float32)
    inverse = tl.load(R + row)
    u = dy * gamma
    mean_ux = tl.sum(u * x, 0) / N
    dx = inverse * u - x * inverse * inverse * inverse * mean_ux
    tl.store(DX + row * N + column, dx, column < N)
    tl.store(PARTIAL + row * N + column, dy * x * inverse, column < N)


def backward_rows(x, gamma, dy, inverse, epsilon):
    check_contract(x, gamma, epsilon)
    if (
        dy.shape != x.shape
        or dy.dtype != x.dtype
        or dy.device != x.device
        or not dy.is_contiguous()
    ):
        raise ValueError("dy must match x and be contiguous")
    m, n = x.shape
    dx = torch.empty_like(x)
    partial = torch.empty((m, n), device=x.device, dtype=torch.float32)
    rmsnorm_dx_kernel[(m,)](
        x, gamma, dy, inverse, dx, partial, n, triton.next_power_of_2(n)
    )
    return dx, partial

Kernel này tạo ra gì

Chương trình tải x, gamma và dy ở FP32, rồi lấy nghịch đảo đã lưu ở forward. mean_ux cộng u × x và một lần nữa chia cho N. Các cột bị che không góp vào tổng này.

Sau đó nó ghi hai đối tượng. DX nhận gradient đầu vào, có cùng dạng và kiểu với x. PARTIAL nhận một đóng góp dgamma cho từng hàng và cột: dy × x × inverse. Bảng này chưa là gradient gamma cuối cùng; nó được reduce ở bước tiếp.

Trong backward_rows(), kiểm tra thêm yêu cầu dy cùng dạng, dtype, thiết bị với x và contiguous. Wrapper Autograd sau đó gọi dy.contiguous() để đáp ứng điều kiện cuối. Điều này không làm x và gamma không contiguous do người dùng đưa vào trở nên hợp lệ: hợp đồng vẫn từ chối chúng.

Nghịch đảo ở đây là trung gian nội bộ do forward_op() tạo rồi giữ cho chính phép tính đó. Đừng hiểu trích đoạn này là API công khai xác thực mọi tensor inverse từ bên ngoài.

Bây giờ kiểm tra dx với tham chiếu:

docker compose run --rm gpu python -m rmsnorm_kernel.cli check-backward --config configs/minimum.yaml

Nếu forward đúng nhưng dx không đúng, hãy kiểm tra đạo hàm: forward đúng không bảo đảm ta đã triển khai đúng lượt về.

Reduce dgamma không có race khi ghi

Một trọng số dùng chung nhận nhiều đóng góp

Mỗi cột dùng cùng trọng số gamma cho mọi hàng. Gradient của nó phải gom đóng góp các hàng. Hãy hình dung nhiều tác vụ mỗi tác vụ tính một phần tổng: ta cộng chúng, không tùy ý giữ phần cuối.

Nếu mọi chương trình GPU cùng sửa một ô không phối hợp, sẽ có race khi ghi. Phép cộng atomic có thể tránh mất đóng góp, nhưng thứ tự cộng đồng thời có thể khác. Trong số thực dấu phẩy động, thứ tự này có thể đổi các chữ số cuối.

Phiên bản tối thiểu chọn tổ chức khác: mỗi hàng ghi vùng partial riêng, rồi lượt Triton thứ hai cộng các đóng góp theo cột. Nó reduce ở FP32 trước khi đổi về dtype của gamma. Cách này tránh tích lũy atomic đồng thời, nhưng không hứa giống từng bit trên mọi phần cứng và phiên bản.

def reduce_gamma_gradient(partial, dtype):
    m, n = partial.shape
    result = torch.empty(n, device=partial.device, dtype=dtype)
    gamma_sum_kernel[(triton.cdiv(n, 32),)](partial, result, m, n, 32, 32)
    return result

result có N phần tử, một gradient mỗi trọng số. triton.cdiv(n, 32) là phép chia làm tròn lên: với N=257, nó chuẩn bị chín chương trình. Thân gamma_sum_kernel không nằm trong trích đoạn; tổ chức đầy đủ của lượt reduce thứ hai ở archive. Hàm trên cho thấy cách chạy và kiểu đầu ra.

Tổng không phải trung bình

Giả sử, trong ví dụ thuần sư phạm, hai hàng đã tạo các đóng góp sau:

Đóng góp	Cột 0	Cột 1
Hàng 0	1	2
Hàng 1	3	4
Tổng mong đợi cho `dgamma`	4	6
Trung bình, sai trong trường hợp này	2	3

Công thức dgamma không có phép chia thêm cho M. Nếu loss toàn cục đã lấy trung bình trên các ví dụ, trọng số đó truyền qua dy. Chia thêm ở đây sẽ đổi gradient phép toán.

Với một hàng, tổng và trung bình giống nhau. Với 17 hàng, chúng thường khác: vì vậy các trường hợp M=1 và M=17 hữu ích để phát hiện nhầm lẫn.

Cái giá của tổ chức này: một bảng tạm

partial chứa M × N số FP32, mỗi số bốn byte. Vì thế dung lượng là M × N × 4 byte.

Với dạng benchmark 128 × 257, nó là 131 584 byte, tức 128,5 Kio. Tính này chỉ nói về partial, không phải toàn bộ VRAM. Với M=4096 và N=4096, bảng này đạt 64 Mio; thêm một hàng sẽ vượt ngưỡng.

Các phép tính này giải thích ngân sách scratch: chiều rộng N được phép không đủ để mọi chiều cao M đều hợp lệ. Chúng không phải phép đo bộ nhớ mới của dự án.

docker compose run --rm gpu python -m rmsnorm_kernel.cli check-gamma --config configs/minimum.yaml

Kiểm tra check-gamma bổ sung các kiểm tra trước: nó nhắm vào gradient trọng số dùng chung, không chỉ gradient đầu vào.

Kết nối Autograd và thực sự huấn luyện

Ta đã có forward, dx và dgamma. Còn phải nối chúng vào PyTorch. Đây là vai trò của lớp con torch.autograd.Function.

Với lập trình viên web, có thể xem nó là adapter: phần còn lại gọi một thao tác bình thường, adapter biết ủy quyền thực thi và lượt về cho mã chuyên biệt. Autograd không đoán công thức Triton: ta cung cấp rõ quy tắc backward. Tài liệu PyTorch về hàm tùy biến mô tả giao diện này.

class RMSNormFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, gamma, epsilon):
        y, inverse = forward_op(x, gamma, epsilon)
        ctx.save_for_backward(x, gamma, inverse)
        ctx.epsilon = epsilon
        CALLS["forward"] += 1
        return y

    @staticmethod
    def backward(ctx, dy):
        x, gamma, inverse = ctx.saved_tensors
        dx, partial = backward_rows(x, gamma, dy.contiguous(), inverse, ctx.epsilon)
        dgamma = reduce_gamma_gradient(partial, gamma.dtype)
        CALLS["backward"] += 1
        return dx, dgamma, None


def rmsnorm(x, gamma, epsilon=1e-6):
    return RMSNormFunction.apply(x, gamma, epsilon)

ctx giữ lại gì

ctx là ngữ cảnh của một lần gọi hàm. Ở forward, save_for_backward(x, gamma, inverse) giữ tensor cần cho lượt về. epsilon, không phải tensor học được trong hợp đồng này, được giữ riêng như thuộc tính.

Ở backward, ta lấy lại các tensor đó. dy.contiguous() bảo đảm bố trí gọn của gradient đi lên: phương thức trả tensor nếu đã phù hợp hoặc tạo bản sao contiguous khi cần, như tài liệu contiguous() nêu.

backward_rows() tính dx và đóng góp trung gian; reduce_gamma_gradient() cộng chúng. Cuối cùng return dx, dgamma, None trả lời ba đối số đầu vào cùng thứ tự: gradient cho x, cho gamma, và không có cho epsilon.

None này không nói epsilon không có tác động toán học lên đầu ra. Nó nói tham số này không được lấy đạo hàm trong giao diện của ta.

Gradient chưa phải cập nhật

Không có lời gọi optimizer nào trong lớp này. Nó tính đầu ra và trả gradient; vòng huấn luyện bên ngoài sau đó dùng Adam để sửa tham số. Bộ đếm CALLS ghi lượt qua wrapper, không phải số lần chạy GPU chính xác: một backward ở đây gọi tính theo hàng rồi reduce gamma.

Mạng nhỏ so sánh là Linear → RMSNorm → SiLU → Linear. Lớp Linear kết hợp đầu vào với trọng số; SiLU thêm biến đổi phi tuyến. Hai nhánh bắt đầu từ khởi tạo giống nhau, nên ta so sánh tham chiếu và kernel mà không cho chúng điểm xuất phát khác.

Lộ trình huấn luyện 100 bước với Adam và 32 ví dụ mỗi bước, nạp lại trọng số, rồi kiểm tra nhánh kernel đã gọi forward và backward 100 lần. Kiểm tra này tránh xác nhận một lần huấn luyện thực ra vẫn đi nhầm theo nhánh tham chiếu.

docker compose run --rm gpu python -m rmsnorm_kernel.cli train-integration --config configs/minimum.yaml

Tích hợp kiểm tra phép toán thực sự tham gia học. Việc nạp lại cũng kiểm tra trọng số đã lưu. Đây không phải cùng khẳng định với việc tiếp tục hoàn toàn sau gián đoạn, khôi phục mọi trạng thái optimizer; kịch bản sau không được nêu ở đây.

Các phép đo trước khi xem kết quả

Ba so sánh số học riêng biệt

Các kiểm tra bao phủ 72 trường hợp mỗi loại cho forward, dx và dgamma. Vì vậy ta so sánh đầu ra, gradient đầu vào và gradient trọng số: ba đối tượng khác nhau.

Dung sai công bố là atol=1e-5, rtol=1e-4 ở FP32 và atol=rtol=3e-3 ở FP16. Atol là biên tuyệt đối; rtol thêm biên tỷ lệ theo giá trị tham chiếu.

Để đọc các tham số này, quy tắc gần nhau của torch.allclose viết cho từng phần tử là:

abs(valeur - référence) ≤ atol + rtol × abs(référence)

Với dung sai FP32, tham chiếu bằng 1 cho phép chênh 0,00011; tham chiếu bằng 100 cho phép 0,01001. Các ví dụ giải thích phép tính, không phải trường hợp ghi lại trong chiến dịch.

Một sai số tuyệt đối tối đa, riêng nó, không đủ kết luận dung sai tương đối đã bị vượt. Cần biết dtype, giá trị tham chiếu và kết quả kiểm tra tương ứng. Ngược lại, trung bình sai số nhỏ không cho phép bỏ qua một phần tử thất bại.

Vì sao không luôn đòi bằng 0? Số dấu phẩy động có độ chính xác hữu hạn và thứ tự thao tác có thể đổi làm tròn, ngay cả khi biểu thức toán học tương đương. Ghi chú PyTorch về độ chính xác số mô tả giới hạn này. Ta giữ dung sai được công bố trước khi so sánh, thay vì tăng chúng sau khi thấy lỗi.

Benchmark đo thời gian nào?

Giao thức đặt sự kiện CUDA quanh các lời gọi hoàn chỉnh. Nó giữ lại cấp phát, đường đi qua Autograd và các lần chạy liên quan; không thu hẹp phép đo chỉ vào thân kernel Triton.

Một điểm cần rõ: sự kiện CUDA đo khoảng trên dòng thời gian GPU. Chúng không phải đồng hồ phân chi phí CPU riêng cho từng dòng Python. Kết quả mô tả phạm vi chạy này, không phải phân rã chi tiết cấp phát, dispatch và tính toán.

Warm-up chạy thao tác trước khi giữ thời lượng. Ở đây loại 20 lời gọi warm-up; trung vị tính trên 50 mẫu, mỗi mẫu 20 lời gọi. Hai số 20 khác vai trò: số đầu để chuẩn bị, số sau để gom lời gọi đã đo.

Trung vị mô tả điểm giữa của quan sát đã sắp thứ tự, khác trung bình là cộng thời lượng rồi chia số lượng. Nó không tự mô tả độ trễ tệ nhất hay toàn bộ biến thiên.

docker compose run --rm gpu python -m rmsnorm_kernel.cli benchmark --config configs/minimum.yaml

Quan sát của lộ trình

Các độ trễ sau là phép đo cục bộ quan sát trên RTX 3090, cho dạng 128 × 257 ở FP32. Ba dòng lỗi mỗi dòng tóm tắt chiến dịch 72 trường hợp; bảng không chia cực đại này theo dtype. Đừng gán chúng cho một trường hợp riêng không nêu trong tóm tắt.

Các số này không dự đoán kernel tốt nhất cho dạng khác, card khác, so sánh với torch.compile hay phạm vi bấm giờ khác.

Trên di động, hãy cuộn ngang bảng.

Kiểm tra quan sát	Kết quả
Sai số tuyệt đối tối đa forward, 72 trường hợp	0,00024414
Sai số tuyệt đối tối đa dx, 72 trường hợp	0,00048828
Sai số tuyệt đối tối đa dgamma, 72 trường hợp	0,00012207
Trung vị tham chiếu forward	52,17 µs
Trung vị kernel forward	70,71 µs
Trung vị tham chiếu forward + backward	267,34 µs
Trung vị kernel forward + backward	251,44 µs

Tích hợp cho sai số đầu ra ban đầu, hay output, là 1,19e-7 và sai số gradient 5,96e-8. Sau 100 bước, nhánh tham chiếu và kernel có MSE test khoảng 0,30670 và lần nạp lại có chênh lệch bằng 0. MSE là trung bình bình phương chênh lệch giữa dự đoán và mục tiêu, không phải phần trăm thành công. Nhánh kernel đã thực thi 100 forward và 100 backward.

Các kiểm tra này áp dụng cho mạng nhỏ và dữ liệu tổng hợp. Chúng không đo cải thiện khả năng hiểu của mô hình ngôn ngữ hay chất lượng nghiệp vụ.

Forward và lộ trình huấn luyện không kể cùng một câu chuyện

Bảng sau chỉ chứa chênh lệch tính từ các trung vị công bố ở trên, không phải phép đo mới. Phần trăm chia chênh lệch thời lượng cho trung vị tham chiếu.

Phạm vi	Chênh lệch kernel − tham chiếu	Biến thiên thời lượng
Chỉ forward	`70,71 − 52,17 = +18,54 µs`	Khoảng +35,54 %
Forward + backward	`251,44 − 267,34 = −15,90 µs`	Khoảng −5,95 %

Vì vậy forward riêng chậm hơn với dạng đã đo này. Cặp forward-backward nhanh hơn trong phạm vi này. Kết quả có thể đúng số học mà không giảm độ trễ trong mọi cách dùng.

Ta không suy ra “thời gian backward thuần” bằng cách trừ hai trung vị của hai lộ trình riêng. Ta cũng không gán chênh lệch cho nguyên nhân cụ thể nếu không profiling thêm. Các số không chi tiết sự gán này hay độ phân tán thời gian đầy đủ.

Với ứng dụng chỉ dự đoán, forward trực tiếp quan trọng. Với vòng học, cả forward và backward đều quan trọng, nhưng chúng vẫn không tự đại diện cho toàn mạng, optimizer và nạp dữ liệu. Kết luận phải gắn với công việc thực sự được bấm giờ.

Điều chúng ta đã học

Ta bắt đầu từ trung bình bình phương và xây dựng đường hoàn chỉnh: tham chiếu dễ đọc, forward Triton, backward đầu vào, reduce trọng số dùng chung và nhánh Autograd dùng khi huấn luyện.

Độ đúng phụ thuộc chi tiết mà vài dòng công thức có thể che giấu. N mô tả dữ liệu; BLOCK mô tả công việc của chương trình GPU. Mặt nạ ngăn truy cập ngoài hàng, nhưng không sửa mẫu số sai. Chuyển sang FP32 bảo vệ phép tính trung gian trong phạm vi chọn; hợp đồng contiguous bảo vệ giả định địa chỉ.

Gradient của gamma cũng buộc ta chọn tổ chức bộ nhớ: đóng góp riêng rồi tổng theo cột. Lựa chọn này tránh tích lũy đồng thời trực tiếp, đổi lại scratch bị giới hạn.

Cuối cùng, mỗi bằng chứng giữ nguyên phạm vi. Gradcheck xác thực tham chiếu float64; so sánh số học đánh giá kernel FP32/FP16; huấn luyện và nạp lại kiểm tra tích hợp. Không kiểm tra nào tự mở rộng hỗ trợ sang dtype, stride hay đạo hàm bậc cao khác, và không kiểm tra nào bảo đảm tăng tốc phổ quát.

Bài tập: giả thuyết trước lệnh

Ta sẽ xem một lỗi rất kín: thay mẫu số N bằng độ rộng block làm tròn. Trong đề, BLOCK_SIZE là độ rộng này, được gọi BLOCK trong các kernel đã trình bày.

Trước khi chạy lệnh, lấy hàng 257 giá trị đều bằng 1. Tổng bình phương là bao nhiêu? Trung bình nhận được khi chia 257 rồi 512? Mẫu số RMSNorm lớn hơn hay nhỏ hơn? Biên độ đầu ra thế nào?

Làm lại lập luận trong đầu với 256 giá trị. Giải thích vì sao test chỉ ở kích thước này có thể bỏ lọt lỗi.

docker compose run --rm gpu python -m rmsnorm_kernel.cli exercise --config configs/exercise.yaml

Bây giờ so sánh kịch bản âm với giả thuyết. Bài tập phải làm lỗi quan sát được; không nhằm thay kernel tham chiếu bằng biến thể sai.

Đây là mốc số học độc lập, với gamma=1 và epsilon=1e-6, để kiểm tra hướng giải thích. Đây không phải đầu ra được chụp từ lệnh:

Hàng 257 giá trị bằng 1	Chia đúng cho N	Chia sai cho BLOCK
Tổng bình phương	257	257
Trung bình dùng	1	`257 / 512 = 0,501953125`
Mẫu số với epsilon	Khoảng `1,0000005`	Khoảng `0,708487`
Đầu ra mỗi tọa độ	Khoảng `0,9999995`	Khoảng `1,411458`

Phép chia sai đánh giá thấp mẫu số và khuếch đại đầu ra quá mức. Số 0 đệm không đổi tổng, nhưng đổi trung bình vì bị tính nhầm. Ở N=256, N và BLOCK trùng nhau: lỗi cụ thể này không còn xuất hiện.

Để hoàn tất bài tập, cũng giải thích vì sao backward phải dùng cùng N logic trong mean(u × x). Chỉ sửa forward không đủ cho đạo hàm đúng.

Điều chỉnh kernel và đưa vào production

Mở rộng hợp đồng trước hết cần test mới. BF16, stride tùy ý, N lớn hơn 4096 và double backward không phải khả năng mà hướng dẫn này đã có. BF16 là định dạng số khác; không thể suy ra hỗ trợ nó từ FP16.

Double backward tính đạo hàm của gradient. Nó không chỉ là làm hai bước huấn luyện. Hàm cho phép backward đầu tiên không tự động dùng được cho đạo hàm bậc cao. Wrapper này không xác thực kịch bản đó.

Để nhận bảng lớn hơn, hãy định nghĩa lại ngân sách scratch và kế hoạch reduce. Để nhận bố trí bộ nhớ khác, hãy sửa địa chỉ theo stride và test. Đừng chỉ xóa từ chối trong check_contract(): việc đó không làm giả định kernel thành đúng.

Fallback là quay lại tham chiếu một cách rõ ràng khi đầu vào không hỗ trợ. Dự án không cung cấp fallback im lặng. Thêm nó cần định nghĩa điều kiện chuyển, test cả hai đường và làm việc chuyển đó quan sát được.

Cuối cùng, đưa production cần các dạng đại diện mạng thật, load test, đo phạm vi cần thiết, tương thích driver/compiler, khả năng quan sát và chiến lược quay về tham chiếu. Corpus tổng hợp và một host không đủ chứng minh các thuộc tính này. So sánh với triển khai khác phải giữ cùng công thức, epsilon và yêu cầu độ chính xác rõ ràng.

Khắc phục sự cố và dọn dẹp

Forward chủ yếu thất bại ở N=257. Kiểm tra chia cho N, mặt nạ column < N và địa chỉ dùng row * N. Kết quả đúng ở 256 không minh oan cho việc nhầm dữ liệu với block công việc.

Forward qua nhưng dx lệch. Kiểm tra hạng r³, trung bình u × x và nghịch đảo giữ từ forward. Đừng vội chỉnh optimizer: đạo hàm sai đã là vấn đề trước huấn luyện.

Chỉ dgamma lệch hoặc thay đổi. Kiểm tra đóng góp partial, lượt reduce thứ hai và ngân sách scratch. Xác nhận đó là tổng theo hàng, không phải trung bình.

Đầu vào không contiguous bị từ chối. Việc từ chối là dự kiến: hợp đồng không có fallback im lặng. Làm dy contiguous trong wrapper không đổi quy tắc cho x và gamma.

Sau khi giữ lại quan sát, hãy tạo hình, kiểm tra checkpoint, xem trước dọn dẹp rồi dừng dịch vụ. Bốn lệnh này kết thúc lộ trình theo thứ tự:

docker compose run --rm gpu python lesson/figures/plot.py
docker compose run --rm gpu python scripts/check_checkpoints.py
docker compose run --rm gpu python scripts/clean.py --scope generated --dry-run
docker compose down --remove-orphans

--dry-run không xóa gì. Hãy đọc tệp chương trình đề nghị dọn trước mọi thao tác phá hủy. Lệnh cuối dừng dịch vụ của dự án Compose này; nó không phải dọn sạch toàn bộ môi trường Docker.

Nguồn gốc

Các tham chiếu dưới đây đào sâu cơ chế. Ví dụ số học trong bài là minh họa sư phạm; phép đo cục bộ vẫn là của lộ trình r1. Liên kết tài liệu theo phiên bản không thay thế phụ thuộc trong ZIP và không buộc đổi phiên bản.

RMSNorm — bài báo gốc : nguyên tắc chuẩn hóa không trừ trung bình.
Triton — hướng dẫn LayerNorm : bối cảnh reduction và backward chuẩn hóa.
PyTorch — Autograd : cơ chế vi phân tự động.
NVIDIA Container Toolkit — hướng dẫn cài đặt : truy cập GPU từ Docker.

Để khám phá lập trình Triton, hãy bắt đầu từ ví dụ cộng vector. Để xem lại hợp đồng toán học, cũng xem RMSNorm trong PyTorch.

Để nối vào mạng, các trang Extending PyTorch, gradcheck và Tensor.contiguous bổ sung cho phần wrapper.

Để diễn giải phép đo, hãy giữ sẵn torch.allclose, ghi chú về độ chính xác số và tài liệu sự kiện CUDA.