Phân tích hiệu năng mô hình PyTorch trên AWS Trainium với TorchNeuron

AWS Trainium là một con chip được thiết kế riêng cho học sâu, nhưng việc tối ưu hóa mô hình cho nó đòi hỏi các kỹ thuật phân tích hiệu năng (profiling) cụ thể. Một hướng dẫn mới từ Hugging Face minh họa cách sử dụng TorchNeuron, backend PyTorch gốc dành cho Trainium, để phân tích hiệu năng suy luận một cách có hệ thống. Hướng dẫn này đưa các kỹ sư qua từng bước: ghi lại các trace, đọc dữ liệu cấp hệ thống và so sánh các chiến lược triển khai khác nhau cho một phép toán phổ biến.

Điều gì đã xảy ra

Nhiều người làm máy học dựa vào trực giác để tăng tốc mô hình, nhưng ít người sở hữu các công cụ để chẩn đoán chính xác lý do tại sao mô hình chạy chậm. Để giải quyết khoảng trống này, Hugging Face đã công bố phần đầu tiên của một loạt bài viết tập trung vào profiling cấp hệ thống cho AWS Trainium. Hướng dẫn giả định người dùng chưa có kinh nghiệm trước đó về profiling trên phần cứng này và bắt đầu từ con số không, sử dụng RMSNorm làm trường hợp nghiên cứu.

Bài walkthrough sử dụng TorchNeuron, một backend mã nguồn mở cho phép các quy trình làm việc PyTorch tiêu chuẩn chạy trên Trainium với những thay đổi mã tối thiểu. Nó hỗ trợ thực thi eager, torch.compile và các toán tử tùy chỉnh thông qua Neuron Kernel Interface. Mục tiêu là dạy các kỹ sư cách cấu hình profiler, ghi lại các trace sạch và diễn giải dữ liệu thu được để xác định các phép toán tốn thời gian.

Kết thúc hướng dẫn, người đọc sẽ biết cách tương quan thời gian CPU với thời gian thiết bị, theo dõi các phép toán xuyên suốt stack phần mềm và hiểu cách các mẫu mã hóa khác nhau ảnh hưởng đến mức độ sử dụng phần cứng. Kiến thức nền tảng này giúp các nhà phát triển sẵn sàng tối ưu hóa các mô hình lớn hơn bằng cách xác định các điểm nóng (hotspots) và áp dụng các kỹ thuật fusion hoặc biên dịch phù hợp.

Cách hoạt động

Profiling trên Trainium phụ thuộc vào hai thành phần chính: NeuronConfig và NeuronProfiler. Đối tượng NeuronConfig kiểm soát các luồng dữ liệu cần ghi lại, chẳng hạn như chỉ thị thiết bị và trace runtime, đồng thời chỉ định thư mục đầu ra. Nó được truyền vào torch.profiler.profile dưới dạng cấu hình thử nghiệm. NeuronProfiler sử dụng cấu hình này để xuất cả trace framework tiêu chuẩn và trace hệ thống đặc thù của Neuron.

Để trực quan hóa dữ liệu, các kỹ sư sử dụng Neuron Explorer, một trình xem dựa trên trình duyệt có thể truy cập qua SSH tunnel hoặc tiện ích mở rộng VS Code. Công cụ này hiển thị trace hệ thống, dòng thời gian thiết bị và phân tích chi tiết theo từng engine. Nó khớp các file thực thi đã biên dịch (NEFF files) với các trace chỉ thị phần cứng tương ứng (NTFF files) bằng cách sử dụng một giá trị băm số chung. Điều này cho phép người dùng thấy chính xác cách các phép toán PyTorch cấp cao ánh xạ sang việc thực thi phần cứng cấp thấp.

Quy trình profiling bao gồm việc đồng bộ hóa thiết bị bên trong vòng lặp profiler để đảm bảo tính chính xác về thời gian. Nếu không gọi torch.neuron.synchronize(), profiler có thể xuất dữ liệu trước khi thiết bị hoàn tất công việc, dẫn đến các cột thời gian trống hoặc sai lệch. Các trace kết quả cho thấy vòng đời của mỗi phép toán từ lúc dispatch trong framework, lên lịch trong runtime và cuối cùng là thực thi trên phần cứng NeuronCore.

Chi tiết chính

• TorchNeuron tích hợp với PyTorch bằng cách sử dụng điểm mở rộng PrivateUse1, tránh nhu cầu phải có một khóa dispatch chuyên biệt như CUDA.
• Tham số correlate_device_time=True của NeuronConfig lan truyền thời gian thiết bị vào bảng key_averages() để so sánh trực tiếp với thời gian CPU.
• Các triển khai ngây thơ (naive implementations) của RMSNorm sử dụng các phép toán PyTorch riêng lẻ dẫn đến nhiều lần dispatch phần cứng tách biệt và tạo ra các tensor trung gian.
• Việc sử dụng torch.rms_norm() hạ cấp phép toán xuống thành một kernel Neuron fused duy nhất, giảm đáng kể thời gian thiết bị bằng cách loại bỏ các khoảng trống giữa các phép toán.
• torch.compile với backend Neuron có thể tự động fuse các phép toán naive thành một đồ thị duy nhất, đạt hiệu năng tương đương với fusion thủ công mà không cần thay đổi cấu trúc mã Python ban đầu.
• Neuron Explorer tổ chức các sự kiện vào các lane cho dispatch framework, tương tác thiết bị, lời gọi API runtime và thực thi phần cứng, cho phép phân tích chi tiết chuỗi phụ thuộc.

Tại sao điều này quan trọng

Đối với các kỹ sư xây dựng sản phẩm trên AWS Trainium, việc hiểu rõ sự khác biệt giữa overhead CPU và tính toán thực tế trên thiết bị là rất quan trọng để tối ưu hóa chi phí và độ trễ. Mã naive thường trông hiệu quả trong Python nhưng lại tạo ra quá nhiều overhead trên bộ tăng tốc do các lần dispatch và chuyển bộ nhớ thường xuyên. Bằng cách trực quan hóa các khoảng trống này, các nhà phát triển có thể đưa ra quyết định sáng suốt về khi nào nên sử dụng các phép toán fused tích hợp sẵn so với việc dựa vào các tối ưu hóa của compiler.

Khả năng theo dấu một phép toán từ framework PyTorch xuống tận phần cứng NeuronCore cung cấp mức độ minh bạch thường thiếu ở các bộ tăng tốc hộp đen. Tầm nhìn này giúp các nhóm gỡ lỗi các hồi quy hiệu năng và xác nhận rằng nỗ lực tối ưu hóa của họ thực sự đang tác động đến phần cứng như mong đợi. Nó chuyển dịch việc tối ưu hóa từ phỏng đoán sang kỹ thuật dựa trên bằng chứng.

Hơn nữa, khi các mô hình trở nên phức tạp hơn, hiệu ứng lũy kế của những kém hiệu quả nhỏ trở nên đáng kể. Việc xác định và fuse các phép toán sớm trong chu kỳ phát triển ngăn ngừa nợ kỹ thuật liên quan đến hiệu năng. Các kỹ thuật được mô tả ở đây cung cấp một quy trình làm việc lặp lại để kiểm toán hiệu quả mô hình, đảm bảo rằng chi tiêu cơ sở hạ tầng phù hợp với nhu cầu tính toán.

Bạn có thể làm gì

• Cài đặt gói torch_neuronx và cấu hình NeuronConfig với ProfileMode.DEVICE và ProfileMode.RUNTIME để ghi lại các trace toàn diện.
• Sử dụng torch.profiler.schedule với các bước warmup để loại trừ overhead khởi tạo khỏi các phép đo hiệu năng của bạn.
• Gọi torch.neuron.synchronize() bên trong vòng lặp profiling để đảm bảo các tác vụ thiết bị hoàn tất trước khi xuất trace.
• Phân tích bảng key_averages() được sắp xếp theo self_privateuse1_time_total để xác định các phép toán đắt đỏ nhất trên thiết bị Neuron.
• Tải các file .neff và .ntff được tạo lên Neuron Explorer để trực quan hóa dòng thời gian thực thi và xác minh việc fusion toán tử.
• So sánh các triển khai naive với các biến thể torch.rms_norm hoặc torch.compile để định lượng lợi ích hiệu năng của việc fusion phép toán.