Triển khai Mamba trong một tệp PyTorch

mamba-minimal

• Giới thiệu về mamba-minimal, một bản triển khai Mamba đơn giản chỉ trong một tệp được viết bằng PyTorch.
• Cho ra kết quả số giống hệt bản triển khai chính thức, áp dụng cho cả lan truyền xuôi và lan truyền ngược.
• Mã nguồn được tinh gọn, dễ đọc và có chú thích.
• Không bao gồm các tính năng như tối ưu hóa tốc độ của bản triển khai chính thức.
• Không bao gồm khởi tạo tham số phù hợp, nhưng có thể bổ sung mà không làm giảm tính dễ đọc.

Demo

• Tệp demo.ipynb cho thấy ví dụ hoàn thành prompt.
• Cung cấp ví dụ sinh văn bản bằng mô hình Mamba và AutoTokenizer.
• Trong văn bản ví dụ được tạo ra, Mamba được mô tả là loài rắn độc dài nhất thế giới.

References

• Kiến trúc Mamba được giới thiệu trong bài báo "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" do Albert Gu và Tri Dao viết.
• Có thể xem bản triển khai chính thức trên GitHub.

Ý kiến của GN⁺

• mamba-minimal là một dự án được tạo ra để giảm bớt độ phức tạp của các bản triển khai Mamba hiện có, giúp ngay cả kỹ sư phần mềm mới bắt đầu cũng có thể hiểu được.
• Dự án này góp phần nâng cao tính dễ đọc và khả năng hiểu mã trong lĩnh vực học máy.
• Thông qua các ví dụ sử dụng thực tế, dự án cho thấy một cách dễ hiểu cách khai thác mô hình Mamba, và đây có thể là tài liệu rất thú vị với người học.

[GN⁺]

Ý kiến Hacker News

• Chia sẻ thư viện

  • Đã tạo một thư viện cùng đồng nghiệp để trích xuất phần lớn mã mô hình dùng chung. Phần lớn các mô hình có thể được triển khai trong khoảng 100 dòng, không tính phần import của Python và chú thích.
  • Lấy BERT, Llama 1/2 và MPT làm ví dụ, đồng thời hỗ trợ TorchScript JIT, PyTorch flash attention, v.v.

• Chia sẻ triển khai suy luận Mamba

  • Mã Mamba có nhiều phần như tối ưu tốc độ nên khó áp dụng trực tiếp vào huấn luyện, vì vậy tác giả chia sẻ phần triển khai suy luận Mamba của mình để hỗ trợ việc này.

• Yêu cầu giải thích Mamba cho người không chuyên

  • Có người đề nghị một lời giải thích dành cho người không chuyên về trực giác cốt lõi của Mamba, mô hình không gian trạng thái, lý do Mamba thành công, và những lợi ích khác của Mamba ngay cả khi không quan tâm đến độ dài ngữ cảnh.

• Kỳ vọng cốt lõi của thuật toán

  • Ban đầu đã kỳ vọng cốt lõi của thuật toán Mamba sẽ là phép quét tiền tố song song, nhưng mã thực tế lại được triển khai theo cách khác.

• Một câu đùa dí dỏm về Mamba

  • Bật cười với câu đùa rằng Mamba là loài rắn độc dài nhất thế giới. Thấy thú vị vì có thể tham khảo bài báo arXiv để hé nhìn những gì ẩn sau nghiên cứu này.

• Câu hỏi về độ khó khi huấn luyện mô hình Mamba

  • Có người hỏi việc huấn luyện mô hình Mamba trên Huggingface khó đến mức nào, cũng như cần bao nhiêu GPU và bao nhiêu thời gian để huấn luyện mô hình lớn nhất.

• Chia sẻ nỗ lực diễn giải phiên bản CUDA chính thức

  • Đã từng cố gắng diễn giải phiên bản CUDA chính thức nhưng thất bại và sau đó không thử lại nữa. So với điều đó, họ đánh giá triển khai mới tốt hơn nhiều.

• Sự thán phục với triển khai PyTorch một tệp duy nhất

  • Bày tỏ sự thán phục với một triển khai PyTorch chỉ trong một tệp nữa, và cho rằng những công việc như vậy rất quan trọng cho nghiên cứu ML hiệu quả. Tốc độ tiến triển nghiên cứu có liên hệ với nghịch đảo của thời gian chạy thí nghiệm, và điều này lại liên quan đến độ phức tạp Kolmogorov của mã. Nhấn mạnh rằng công cụ nghiên cứu quan trọng đến mức nào trong việc tăng tốc quá trình khám phá tri thức.

• Yêu cầu thảo luận về bài báo gốc

  • Hỏi liệu đã có thảo luận về bài báo gốc hay chưa, đồng thời chỉ ra rằng các mô hình trước đó như RWKV và RetNet đã thiếu kết quả đầy đủ cho độ dài ngữ cảnh 8k do thiếu triển khai hiệu quả, dẫn đến vấn đề bộ nhớ hoặc yêu cầu tính toán không thực tế.

• Lời khen cho việc giản lược nội dung cốt lõi

  • Thích cách những thứ phức tạp được giản lược và trình bày về phần bản chất nhất.