Việc chuyển đổi mô hình Gemma-4 31B của Google sang AWS Inferentia2 inf2.24xlarge đã tạo ra kết quả khớp hoàn toàn từng token so với bản tham chiếu trên CPU—tuy nhiên, mọi câu được tạo ra đều là vô nghĩa. Khoảng cách giữa việc "khớp" (matching) và "hoạt động được" (working) giờ đây đóng vai trò như một lời cảnh báo cho bất kỳ ai đang cố gắng đưa các mô hình LLM khổng lồ lên các chip suy luận tùy chỉnh của Amazon.

Tại sao khớp từng token là chưa đủ

Nhà phát triển đã so sánh từng token đầu ra từ thiết bị Inferentia với token được tạo ra bởi bản chạy mô hình trên CPU. Các luồng dữ liệu hoàn toàn giống nhau, vì vậy phần cứng có vẻ như đã tái tạo chính xác bản thực thi tham chiếu. Trên thực tế, cả hai luồng đều đưa một prompt bị lỗi vào một mô hình đã bị lược bỏ chat template và được cung cấp các dấu hiệu chuyển lượt (turn markers) sai. Việc thiếu template đã khiến mô hình rơi vào một vòng lặp vô tận, tạo ra những nội dung vô nghĩa. Phần cứng đã làm tốt việc của nó—nó đã tái tạo lại một lỗi tồn tại trong mã tham chiếu.

Bài học rất đơn giản: SEQ_MATCH (sự bằng nhau của các token tuần tự) không đồng nghĩa với tính chính xác. Nếu bản thực thi tham chiếu bị lỗi, một bản sao phần cứng trung thành cũng sẽ thừa hưởng cùng một lỗi đó. Việc xác thực phải vượt xa sự tương đồng ở cấp độ token; nó cần các kiểm tra chức năng đầu-cuối (end-to-end) với các đầu vào được định dạng đúng cách.

Các buffer giả dạng tham số

Trong giai đoạn tải, trình tải mô hình đã bỏ qua một thành phần gọi là layer_scalar. Mã nguồn đăng ký đối tượng này dưới dạng một buffer thay vì một parameter trong định nghĩa mô hình PyTorch. Buffer là các tensor tĩnh mà quá trình huấn luyện không cập nhật, và nhiều trình tải thường bỏ qua chúng khi chuyển đổi sang các định dạng tương thích với Neuron. Việc bỏ qua nó đã để các hệ số tỷ lệ (scaling factors) của một số lớp ở giá trị mặc định, làm sai lệch các phép toán trên toàn bộ mạng lưới. Không có lỗi nào được báo lên; mô hình vẫn biên dịch, và pipeline suy luận vẫn chạy, nhưng kết quả số học đã bị sai lệch.

Đối với bất kỳ ai đang chuyển các mô hình lớn sang Inferentia, hãy kiểm tra kỹ mọi tensor không phải là tham số. Ngay cả khi một tensor không được dùng để học, nó vẫn có thể đóng vai trò thiết yếu cho việc tính toán forward-pass chính xác. Việc xác minh thủ công sự hiện diện của buffer có thể ngăn chặn các lỗi tỷ lệ âm thầm mà nếu không sẽ rất khó chẩn đoán.

Sự biến động của spot-instance và quá trình biên dịch kéo dài 39 phút

Chạy một mô hình 31 tỷ tham số trên một spot instance trông có vẻ rẻ, nhưng sự tiết kiệm này đi kèm với các sự kiện thu hồi (reclaim) không thể dự đoán trước. Thời gian biên dịch của nhà phát triển—khoảng 39 phút để chuyển đổi mô hình thành mã tương thích với Neuron—đã biến mất khi AWS thu hồi instance đó. Để sống sót qua các lần gián đoạn, họ đã xây dựng một mạng lưới an toàn ba lớp:

  • ModelBuilder giữ mức sử dụng bộ nhớ trong giới hạn 384 GB của host, tránh các lỗi crash buộc phải khởi động lại.
  • Sao lưu tức thì lên S3 (Immediate S3 mirroring) cho cả các tệp trọng số thô và các tệp “neffs” (tệp thực thi Neuron) đã biên dịch, cho phép một instance mới tiếp tục chính xác tại nơi instance trước đó đã dừng lại.
  • Một bộ quét đa vùng (multi-region poller) quét các vùng AWS để tìm dung lượng spot còn trống và khởi chạy một instance mới ngay khi có thể.

Những bước này đã biến một quá trình biên dịch đơn lẻ, mong manh thành một pipeline kiên cố, có thể tồn tại qua sự biến động của thị trường spot.

Cạm bẫy sharding với các bố cục attention hỗn hợp

Gemma-4 31B sử dụng hai cấu hình attention. Một số lớp sử dụng bốn đầu key-value (KV heads), trong khi các lớp khác sử dụng số lượng khác. Việc chia đều mô hình trên tám rank song song sẽ thất bại khi số lượng KV head của một lớp không chia hết một cách gọn gàng. Việc cố gắng shard một lớp có 4 head trên tám rank sẽ buộc mỗi rank phải xử lý nửa một head—một điều bất khả thi về mặt toán học dẫn đến lỗi sai lệch hình dạng (shape mismatches) và lỗi runtime.

Giải pháp là sao chép các lớp được shard toàn cục (những lớp có số lượng head tương thích) trên tất cả các rank và chỉ shard các lớp “sliding” có số lượng head cho phép chia đều. Chiến lược hybrid này giúp duy trì hiệu suất tensor-parallel trong khi tránh việc chia KV head không hợp lệ, loại bỏ các lỗi tensor-parallelization đã gây khó khăn cho các nỗ lực trước đó.

Lời kết

Việc chuyển đổi một LLM khổng lồ sang Inferentia không chỉ đơn thuần là bài tập biên dịch và chạy. Nó đòi hỏi việc kiểm thử chức năng nghiêm ngặt vượt ra ngoài sự bằng nhau của token, xác minh tỉ mỉ rằng mọi tensor—dù là tham số hay buffer—đều được xử lý chính xác, và một chiến lược triển khai có tính đến việc thu hồi spot-instance. Cuối cùng, việc sharding phải tôn trọng cấu trúc hình học attention nội bộ của mô hình; nếu không, tính song song vốn hứa hẹn tốc độ sẽ trở thành nguồn cơn của những lỗi âm thầm.