Hướng dẫn profiling TorchNeuron trên AWS Trainium để tìm nguyên nhân mô hình chạy chậm
AWS và Hugging Face đã công bố hướng dẫn nhập môn về cách profiling các workload suy luận PyTorch trên AWS Trainium bằng TorchNeuron, backend PyTorch mã nguồn mở dành cho chip Neuron. Phần đầu giải thích cách cấu hình torch.profiler để thu thập hoạt động của CPU, runtime và thiết bị Neuron, sau đó xuất trace để phân tích. Bài viết dùng ba cách triển khai RMSNorm—một phép chuẩn hóa phổ biến trong mạng nơ-ron—gồm các phép toán PyTorch riêng lẻ, kernel fused torch.rmsnorm() và cùng đoạn mã cơ bản được tối ưu bằng torch.compile. Hướng dẫn chỉ cách so sánh thời gian chạy trên thiết bị với thời gian CPU trong bảng keyaverages(), rồi theo dõi một phép toán từ lúc framework gọi lệnh, runtime lập lịch đến khi thực thi trên NeuronCore bằng Neuron Explorer trên trình duyệt. Bài viết cũng giải thích cặp tệp .neff và .ntff, lần lượt chứa chương trình đã biên dịch và trace lệnh phần cứng, đồng thời cung cấp mã để tái hiện trong thư mục ví dụ của torch-neuronx. Bài học thực tế là profiling giúp xác định mô hình chậm vì overhead của framework, khâu lập lịch runtime hay việc thực thi trên phần cứng, từ đó nhà phát triển có thể chọn fusion, biên dịch hoặc kernel Neuron tùy chỉnh thay vì tối ưu theo phỏng đoán.
Các nhà phát triển PyTorch chạy mô hình trên AWS Trainium có thể dùng hướng dẫn và mã mẫu này để đo thời gian CPU, runtime và Neuron, sau đó xác định chính xác khâu gây chậm. Đây là tài liệu hữu ích trước khi quyết định dùng toán tử fused, `torch.compile` hoặc tối ưu kernel tùy chỉnh.
Đọc bài gốc