AttentionGrid
0.0.22
AtightGrid นำเสนอให้คุณโดย Agora เราเป็นองค์กรวิจัย AI แบบโอเพนซอร์สโอเพ่นซอร์สใหม่ทั้งหมดที่อุทิศตนเพื่อพัฒนามนุษยชาติ
เข้าร่วมกับเราที่นี่เพื่อสนับสนุนโครงการนี้หรือรับการสนับสนุน!
AttentionGrid เป็นกรอบการทำงานที่ทันสมัยออกแบบมาเพื่อให้เป็นประชาธิปไตยในการรวมตัวกันของกลไกความสนใจขั้นสูงเข้ากับโมเดล AI ขับเคลื่อนด้วยการพัฒนาล่าสุดในแบบจำลองหม้อแปลงที่มีความสนใจความสนใจเปิดตัวโลกของกลไกความสนใจไปยังผู้ฝึกการเรียนรู้ของเครื่องนักวิจัยและผู้ที่ชื่นชอบ
หากต้องการระเบิดด้วยความสนใจติดตั้งแพ็คเกจโดยใช้ PIP:
pip install AttentionGridการใช้กลไกความสนใจหรือแบบจำลองหม้อแปลงที่มีความสนใจเป็นเรื่องง่ายเหมือน:
from AttentionGrid import BlockwiseParallelJax
import jax . numpy as jnp
# Initialize the class
bpjax = BlockwiseParallelJax (
q_chunk_size = 64 ,
k_chunk_size = 64 ,
hidden_size = 768 ,
num_heads = 12 ,
rotary_dim = 32 ,
intermediate_size = 3072
)
# Suppose we have hidden_states, attention_mask, and position_ids as input data
hidden_states = jnp . random . rand ( 1 , 100 , 768 )
attention_mask = jnp . random . rand ( 1 , 1 , 100 , 100 )
position_ids = jnp . arange ( 100 ). reshape ( 1 , 100 )
# You can now apply the attention mechanism to your input data
output = bpjax . forward ( hidden_states , attention_mask , position_ids )เราขอแนะนำให้คุณแบ่งปันความสนใจกับชุมชนของคุณ! นี่คือลิงค์แชร์อย่างรวดเร็วสำหรับแพลตฟอร์มโซเชียลมีเดียหลายแห่ง:
แบ่งปันบน Twitter
แบ่งปันบน LinkedIn
แบ่งปันบน Facebook
แบ่งปันบน reddit
แบ่งปันบน whatsapp
ขอบคุณที่สนับสนุนความสนใจและมีส่วนร่วมในการทำให้เป็นประชาธิปไตยของ AI! เราสามารถผลักดันขอบเขตของสิ่งที่เป็นไปได้
ในภูมิทัศน์อันกว้างใหญ่ของ AI กลไกความสนใจได้ปฏิวัติความสามารถของเราในการสร้างแบบจำลองที่ทรงพลังซึ่งสามารถมองเห็นรายละเอียดปลีกย่อยในข้อมูลโดยมุ่งเน้นไปที่แง่มุมที่สำคัญและปรับปรุงประสิทธิภาพโดยรวม วิสัยทัศน์ของเรากับ AttentionGrid คือการเชื่อมช่องว่างระหว่างกลไกที่ทันสมัยและแอปพลิเคชันที่ใช้งานได้จริงซึ่งเป็นเครื่องมือที่ทำให้เทคนิคเหล่านี้สามารถเข้าถึงได้และใช้งานง่ายในแอพพลิเคชั่น AI ที่หลากหลาย
AttentionGrid ได้รับการออกแบบด้วยสถาปัตยกรรมที่ใช้งานง่ายและยืดหยุ่นแบ่งออกเป็นสี่องค์ประกอบหลัก:
Core : นี่คือข้อเท็จจริงของกรอบงานของเราคลาสนามธรรมที่อยู่อาศัยซึ่งจัดวางโครงสร้างพื้นฐานสำหรับกลไกความสนใจและแบบจำลองหม้อแปลง
ความสนใจ ?: ไดเรกทอรีที่อุทิศให้กับกลไกความสนใจต่าง ๆ กลไกการให้ความสนใจแต่ละครั้งจะถูกนำไปใช้ตามพิมพ์เขียวที่มีให้ในแกนกลาง
Transformers ?: นี่คือที่โมเดลของหม้อแปลงมามีชีวิตแต่ละรูปแกะสลักตามการออกแบบที่กำหนดไว้ในแกนกลาง
UTILS : กล่องเครื่องมือที่เต็มไปด้วยคลาสผู้ช่วยสำหรับงานที่จำเป็นเช่นการโหลดรุ่นการประมวลผลข้อมูลล่วงหน้าและอื่น ๆ
ตัวอย่าง : demystifying การใช้งานด้วยตัวอย่างมือและสถานการณ์การใช้งาน
โครงสร้างแบบแยกส่วน : ผสมและจับคู่กลไกความสนใจที่แตกต่างกับแบบจำลองหม้อแปลงหลากหลาย
ผู้ใช้เป็นมิตร : เอกสารที่ชัดเจนและตัวอย่างเพื่อช่วยให้คุณเริ่มต้นได้อย่างรวดเร็ว
โอเพ่นซอร์ส : เปิดรับการบริจาค, ความสนใจเติบโตขึ้นจากความรู้ร่วมกันและความคืบหน้าร่วมกัน
สำหรับตัวอย่างรายละเอียดเพิ่มเติมโปรดดูที่โฟลเดอร์ 'ตัวอย่าง' ในที่เก็บของเรา
เราขอเชิญชวนให้มีส่วนร่วมอย่างเปิดเผย! ไม่ว่าคุณจะมีข้อเสนอแนะคุณสมบัติใหม่รายงานข้อผิดพลาดหรือต้องการเพิ่มลงในรหัสของเราโปรดเปิดปัญหาหรือส่งคำขอดึง
AttentionGrid เป็นซอฟต์แวร์โอเพ่นซอร์สที่ได้รับอนุญาตภายใต้ใบอนุญาต Apache
กลไกความสนใจได้เปลี่ยน AI ทำให้เครื่องสามารถ 'โฟกัส' ในส่วนที่สำคัญของข้อมูลอินพุต ด้วยความสนใจเรามุ่งมั่นที่จะทำให้การเข้าถึงเครื่องมืออันทรงพลังเหล่านี้เป็นประชาธิปไตย เราเชื่อว่าอนาคตของ AI นั้นอยู่ในพลังแห่งความสนใจและผ่านความสนใจเราหวังว่าจะเร่งการเดินทางครั้งนี้ สำรวจพื้นที่เก็บข้อมูลของเราเข้าร่วมสาเหตุของเราและนำทางภูมิทัศน์ที่น่าตื่นเต้นนี้ด้วยกัน!
"รายละเอียดไม่ใช่รายละเอียดพวกเขาทำการออกแบบ" - Charles Eames
รวมความสนใจของแฟลชและตัวแปร
บูรณาการสถานที่สำคัญ
รวมความสนใจแบบคู่ขนาน blockwise
รวมความสนใจแฟลชแบบเบาบางแบบไดนามิก
รวมความสนใจข้ามจาก ImageBind
รวมความสนใจของ Colt-5
รวมความสนใจหลายคำถาม
รวม wrappers จาก Lucid Rains X_Transformers, Decoder, Attention, Encoder, Transformer Wrapper
| กลไก | วิธีการโทร | ตัวอย่างนำเข้า |
|---|---|---|
| ความสนใจตนเอง | from AttentionGrid import SelfAttention | from AttentionGrid import SelfAttention |
| ความสนใจระดับโลก | from AttentionGrid import GlobalAttention | from AttentionGrid import GlobalAttention |
| ความสนใจในท้องถิ่น | from AttentionGrid import LocalAttention | from AttentionGrid import LocalAttention |
| ความสนใจแบบลำดับชั้น | from AttentionGrid import HierarchicalAttention | from AttentionGrid import HierarchicalAttention |
| ความสนใจแบบเบาบาง | from AttentionGrid import dynamic_sparse_attention | from AttentionGrid import dynamic_sparse_attention |
| ฟังก์ชั่นขนาดกะทัดรัด | from AttentionGrid import compact | from AttentionGrid import compact |
| ฟังก์ชันดัชนีแผ่น | from AttentionGrid import pad_index | from AttentionGrid import pad_index |
| ความสนใจแบบขนาน | from AttentionGrid import BlockwiseParallelJax | from AttentionGrid import BlockwiseParallelJax |
| ความสนใจแฟลช | from AttentionGrid import FlashAttention | from AttentionGrid import FlashAttention |
| ความสนใจในสถานที่สำคัญ | from AttentionGrid import LandmarkAttention | from AttentionGrid import LandmarkAttention |
| Colt-5 ความสนใจ | from AttentionGrid import Colt5Attention | from AttentionGrid import Colt5Attention |
| ความสนใจหลายครั้ง | from AttentionGrid import MultiQueryAttention | from AttentionGrid import MultiQueryAttention |
| ขยายความสนใจ | from AttentionGrid import DilatedAttention | from AttentionGrid import DilatedAttention |
ฟังก์ชั่น dynamic_sparse_attention ของ Agora ช่วยให้มีความยืดหยุ่นในการเลือกระหว่างการใช้งานแฮช-สปาร์สและการใช้งาน QK-sparse วัตถุประสงค์ของฟังก์ชั่นนี้คือการควบคุมกลไกความสนใจแบบเบาบางแบบไดนามิกโดยใช้ sparsity_mode ที่เลือก
พารามิเตอร์ฟังก์ชันมีดังนี้:
q : Query Tensor of Shape (batch, n_ctx_q, h, d_head)k : คีย์เทนเซอร์ของรูปร่าง (แบทช์, n_ctx_kv, h, d_head)v : ค่าเทนเซอร์ของรูปร่าง (แบทช์, n_ctx_kv, h, d_head)q_idx & k_idx : แสดงดัชนีถังถ้า sparsity_mode คือ 'แฮช' หรือไม่ว่าจะรักษาหัวที่กำหนดถ้า sparsity_mode คือ 'qk' รูปร่างเทนเซอร์คือ (แบทช์, n_ctx_q, h) และ (แบทช์, n_ctx_kv, h) ตามลำดับsm_scale : ค่าคงที่การทำให้เป็นมาตรฐาน, 1/sqrt (d_head) เว้นแต่จะระบุsparsity_mode : 'แฮช' เพื่อเลือกการใช้งานแฮช-สปาร์สและ 'qk' สำหรับการใช้งาน QK-sparse sm_scale คำนวณโดยค่าเริ่มต้นหากไม่ได้ให้ไว้และหากได้รับ sparsity_mode ที่ไม่รู้จักจะได้รับ KeyError
จากนั้นฟังก์ชั่นจะตรวจสอบ sparsity_mode และขึ้นอยู่กับค่าของมันเรียกว่า hash_sparse_attention หรือ qk_sparse_attention
ฟังก์ชั่น compact สร้างการแสดงขนาดกะทัดรัดของอินพุตเทนเซอร์ x โดยใช้ข้อมูลจาก keep_tensor
พารามิเตอร์ฟังก์ชันคือ:
x : อินพุตเทนเซอร์ถึงกะทัดรัดด้วยรูปร่าง (แบทช์, n_ctx, h, d_head)keep_tensor : เทนเซอร์ลอยของรูปร่าง (แบทช์, n_ctx, h) ที่มี 1 เมื่อหัวถูกเก็บไว้และ 0 อย่างอื่น ฟังก์ชั่นแรกคำนวณ indices_per_head ซึ่งคำนวณจำนวนองค์ประกอบที่ไม่ได้ถูกฆ่าต่อหัว มันจัดเรียง keep_tensor ตามลำดับจากมากไปน้อยในขณะที่รักษาลำดับขององค์ประกอบที่เท่ากัน (เสถียร = จริง) จากนั้นรวบรวมองค์ประกอบของ x ตามดัชนีเทนเซอร์ ผลที่ได้คือการเป็นตัวแทนขนาดกะทัดรัดของ x พร้อมกับดัชนีเทนเซอร์และเทนเซอร์ที่แสดงถึงจำนวนองค์ประกอบที่ไม่ได้ถูกฆ่าต่อหัว
ฟังก์ชั่น pad_index แผ่นดัชนีเทนเซอร์เพื่อให้สอดคล้องกับเคอร์เนล ต้องใช้พารามิเตอร์ต่อไปนี้:
index : ดัชนีเทนเซอร์ดั้งเดิมที่ได้รับจาก compact โดยมีรูปร่าง (แบทช์, buffer_size, h) สำหรับแต่ละชุดและการจับเวลาจะแสดงถึงดัชนีหัวที่มีต้นกำเนิดมาจากindices_per_head : สำหรับแต่ละหัวมีจำนวนดัชนีที่ยังไม่ได้ทิ้ง มันสร้างสำเนาของดัชนีเทนเซอร์และสร้างหน้ากากตามขนาดของ indices_per_head จากนั้นจะปรับเปลี่ยนดัชนีในสำเนาที่สอดคล้องกับความจริงในหน้ากากให้เท่ากับ pad_idx
ฟังก์ชั่น qk_sparse_attention เป็นส่วนหนึ่งของกลไกความสนใจแบบกระจายแบบไดนามิก มันถูกใช้เมื่อ sparsity_mode ถูกตั้งค่าเป็น 'qk' ฟังก์ชั่นนี้ใช้กลไกความสนใจของ QK-sparse และต้องการให้พารามิเตอร์ q_keep และ k_keep เป็นประเภทลอย
มันเป็นครั้งแรกที่สร้างการแสดงขนาดกะทัดรัดของเคียวรีคีย์และเทนเซอร์ค่าโดยใช้ฟังก์ชั่น compact จากนั้นแผ่นรองดัชนีเทนเซอร์โดยใช้ฟังก์ชัน pad_index เทนเซอร์จะถูกเปลี่ยนเพื่อเข้ากันได้กับเคอร์เนล ในที่สุดฟังก์ชั่นเรียกฟังก์ชั่น qk_sparse_attention_kernel และทำให้เทนเซอร์ผลลัพธ์กลับเข้าไปในพื้นที่มิติเดิม
ฟังก์ชั่น hash_sparse_attention เป็นส่วนหนึ่งของกลไกความสนใจแบบเบาบางแบบไดนามิก มันถูกใช้เมื่อ sparsity_mode ถูกตั้งค่าเป็น 'แฮช' ฟังก์ชั่นนี้ใช้กลไกความสนใจแฮชสปาร์ส
ฟังก์ชั่นใช้พารามิเตอร์อินพุตเดียวกันกับ qk_sparse_attention อย่างไรก็ตามแทนที่จะเป็นพารามิเตอร์ q_keep และ k_keep ฟังก์ชัน hash_sparse_attention ต้องใช้ q_bucket_idx และ k_bucket_idx ซึ่งเป็นตัวแทนดัชนีถังสำหรับการสืบค้นและคีย์ตามลำดับ
ฟังก์ชั่น hash_sparse_attention เรียงลำดับการสืบค้นคีย์และเทนเซอร์ค่าตามดัชนีถังโดยใช้ฟังก์ชัน sort_bucketed_attention จากนั้นมันจะสร้างการเป็นตัวแทนขนาดกะทัดรัดของเคียวรีคีย์และค่าเทนเซอร์ที่เรียงลำดับโดยใช้ฟังก์ชัน compact จากนั้นแผ่นรองดัชนีเทนเซอร์โดยใช้ฟังก์ชัน pad_index
เทนเซอร์จะถูกเปลี่ยนเพื่อเข้ากันได้กับเคอร์เนล ฟังก์ชั่นนั้นเรียกฟังก์ชั่น hash_sparse_attention_kernel และทำให้เทนเซอร์ผลลัพธ์กลับเข้าไปในพื้นที่มิติเดิม
ฟังก์ชั่น sort_bucketed_attention เป็นฟังก์ชั่นผู้ช่วยที่ใช้ใน hash_sparse_attention มันเรียงลำดับเทนเซอร์อินพุตตามดัชนีถังที่กำหนด
พารามิเตอร์ฟังก์ชันคือ:
qkv : Query, Key, Value Tensors of Shape (Batch, N_CTX, H, D_HEAD)qkv_bucket_idx : ดัชนีถังสำหรับการสืบค้นปุ่มและค่าของรูปร่าง (แบทช์, n_ctx, h) ฟังก์ชั่นแรกเรียงลำดับ qkv_bucket_idx เทนเซอร์และได้รับดัชนีเรียงลำดับ จากนั้นมันก็เรียงลำดับเทนเซอร์ qkv โดยใช้ดัชนีที่เรียงลำดับ นอกจากนี้ยังขยาย qkv_bucket_idx ให้เป็นรูปร่างเดียวกันกับ qkv สำหรับความเข้ากันได้
ฟังก์ชั่น qk_sparse_attention_kernel เป็นฟังก์ชันเคอร์เนลที่ใช้ใน qk_sparse_attention มันคำนวณผลรวมของค่าถ่วงน้ำหนักตาม softmax ของแบบสอบถามและผลิตภัณฑ์หลัก
พารามิเตอร์ฟังก์ชันคือ:
q : Query Tensor of Shape (batch, n_ctx_q, h, d_head)k : คีย์เทนเซอร์ของรูปร่าง (แบทช์, n_ctx_kv, h, d_head)v : ค่าเทนเซอร์ของรูปร่าง (แบทช์, n_ctx_kv, h, d_head)sm_scale : ค่าคงที่การทำให้เป็นมาตรฐาน, 1/sqrt (d_head) เว้นแต่จะระบุ ฟังก์ชั่น hash_sparse_attention_kernel เป็นฟังก์ชันเคอร์เนลที่ใช้ใน hash_sparse_attention มันทำงานได้เช่นเดียวกับ qk_sparse_attention_kernel แต่จัดการกับการจับความสนใจแฮช-สปาร์ส
พารามิเตอร์ฟังก์ชั่นนั้นเหมือนกับของ qk_sparse_attention_kernel อย่างไรก็ตาม q , k และ v ได้รับการจัดเรียงและอัดแน่นขึ้นอยู่กับดัชนีถัง
เคอร์เนลคำนวณผลิตภัณฑ์ของแบบสอบถามและคีย์ปรับขนาดโดย sm_scale ใช้ softmax เพื่อรับน้ำหนักจากนั้นคำนวณผลรวมถ่วงน้ำหนักของค่า
โปรดทราบว่านี่เป็นการตีความทั่วไปของเอกสารและความเข้าใจและการปรับเปลี่ยนฟังก์ชั่นเหล่านี้ในทางปฏิบัติอาจต้องใช้ความรู้เชิงลึกเกี่ยวกับกลไกความสนใจที่กระจัดกระจายและหลักการเรียนรู้เชิงลึก
ฟังก์ ชั่น blockwise_compute_attn :
ฟังก์ชั่น blockwise_compute_attn เป็นส่วนสำคัญของคลาส BlockwiseParallelJax และใช้ในการคำนวณกลไกความสนใจของแบบจำลองในลักษณะ blockwise
พารามิเตอร์:
query , key , value : พารามิเตอร์เหล่านี้เป็นอินพุตหลักสำหรับการคำนวณความสนใจซึ่งแสดงถึงการสืบค้นคีย์และค่าตามลำดับbias : พารามิเตอร์เสริมที่ใช้ในการเพิ่มอคติให้กับคะแนนความสนใจก่อน Softmaxdeterministic : ธงบูลีนที่ใช้ในการตัดสินใจว่าจะใช้การออกกลางคันหรือไม่dropout_rng : เครื่องกำเนิดหมายเลขสุ่มสำหรับการออกกลางคันattn_pdrop : ความน่าจะเป็นของการออกกลางคันเพื่อความสนใจcausal_mask : ธงบูลีนไม่ว่าจะใช้หน้ากากความสนใจเชิงสาเหตุหรือไม่query_chunk_size , key_chunk_size : ขนาดของแต่ละคำถามและคีย์ก้อนตามลำดับdtype : ประเภทข้อมูลของการคำนวณ เป็นค่าเริ่มต้นคือ jnp.float32policy : พารามิเตอร์นี้กำหนดนโยบายสำหรับจุดตรวจสอบการไล่ระดับสีprecision : พารามิเตอร์นี้ใช้เพื่อตั้งค่าระดับความแม่นยำสำหรับการคำนวณ ค่าเริ่มต้นคือ lax.Precision.HIGHESTprevent_cse : ธงบูลีนที่ใช้เพื่อป้องกันการกำจัด subexpression ทั่วไปฟังก์ ชั่น blockwise_compute_ffn :
ฟังก์ชั่น blockwise_compute_ffn ใช้เพื่อคำนวณเครือข่ายฟีดไปข้างหน้าของโมเดลในลักษณะ blockwise
พารามิเตอร์:
cell : เซลล์ในเครือข่ายที่ใช้ฟังก์ชันinputs : ข้อมูลอินพุตสำหรับเครือข่ายฟีดไปข้างหน้าchunk_size : ขนาดของแต่ละก้อนสำหรับการคำนวณ blockwisedeterministic : ธงบูลีนที่ใช้ในการตัดสินใจว่าจะใช้การออกกลางคันหรือไม่policy : พารามิเตอร์นี้กำหนดนโยบายสำหรับจุดตรวจสอบการไล่ระดับสีprevent_cse : ธงบูลีนที่ใช้เพื่อป้องกันการกำจัด subexpression ทั่วไปคลาส Blockwise_lm_head :
คลาส Blockwise_LM_Head เป็นโมดูลที่ใช้การแปลงเชิงเส้นตามด้วยฟังก์ชั่น softmax เพื่อสร้างการกระจายผ่านคำศัพท์สำหรับแต่ละตำแหน่งในอินพุต
vocab_size : ขนาดของคำศัพท์ซึ่งเป็นขนาดของมิติเอาต์พุตของการแปลงเชิงเส้นchunk_size : ขนาดของแต่ละก้อนสำหรับการคำนวณ blockwisepolicy : พารามิเตอร์นี้กำหนดนโยบายสำหรับจุดตรวจสอบการไล่ระดับสีdtype : ประเภทข้อมูลของการคำนวณ เป็นค่าเริ่มต้นคือ jnp.float32prevent_cse : ธงบูลีนที่ใช้เพื่อป้องกันการกำจัด subexpression ทั่วไปฟังก์ ชั่น blockwise_cross_entropy :
ฟังก์ชั่น blockwise_cross_entropy คำนวณการสูญเสีย cross-entropy สำหรับการทำนายของโมเดลในลักษณะ blockwise
พารามิเตอร์:
logits : การทำนายเอาต์พุตของโมเดลtokens : ฉลากที่แท้จริงvalid : หน้ากากที่ระบุตำแหน่งที่ถูกต้องในอินพุตchunk_size : ขนาดของแต่ละก้อนสำหรับการคำนวณ blockwisepolicy : พารามิเตอร์นี้กำหนดนโยบายสำหรับจุดตรวจสอบการไล่ระดับสีprevent_cse : ธงบูลีนที่ใช้เพื่อป้องกันการกำจัด subexpression ทั่วไปคลาส BlockWiseparallelJax :
BlockwiseParallelJax ( q_chunk_size , k_chunk_size , hidden_size , num_heads , rotary_dim , intermediate_size , layer_norm_epsilon = 1e-5 , activation_function = "gelu" , attn_pdrop = 0.0 , resid_pdrop = 0.0 , max_position_embeddings = 1024 , dtype = jnp . float32 , causal = True , policy = 'nothing_saveable' , prevent_cse = False , float32_logits = False )พารามิเตอร์
q_chunk_size : จำนวนเต็ม ขนาดก้อนสำหรับการสืบค้นในความตั้งใจของตนเองk_chunk_size : จำนวนเต็ม ขนาดก้อนสำหรับกุญแจสำคัญในการดูแลตนเองhidden_size : จำนวนเต็ม มิติของชั้นที่ซ่อนอยู่ในหม้อแปลงnum_heads : จำนวนเต็ม จำนวนหัวความสนใจในกลไกการดูแลตนเองrotary_dim : จำนวนเต็มหรือไม่มี จำนวนขนาดที่จะใช้สำหรับการเข้ารหัสตำแหน่งโรตารี่intermediate_size : จำนวนเต็ม ขนาดของชั้นกลางในเครือข่ายฟีดไปข้างหน้าlayer_norm_epsilon : ลอย ค่าคงที่เล็กน้อยเพื่อป้องกันการแบ่งเป็นศูนย์ในการทำให้เป็นมาตรฐานของชั้น ค่าเริ่มต้นคือ 1e-5activation_function : String ฟังก์ชั่นการเปิดใช้งานที่จะใช้ในเครือข่ายฟีดไปข้างหน้า ค่าเริ่มต้นคือ 'gelu'attn_pdrop : ลอย ความน่าจะเป็นแบบกลางคันสำหรับกลไกความสนใจ ค่าเริ่มต้นคือ 0.0resid_pdrop : ลอย ความน่าจะเป็นแบบกลางคันสำหรับการเชื่อมต่อที่เหลือ ค่าเริ่มต้นคือ 0.0max_position_embeddings : จำนวนเต็ม จำนวนสูงสุดของการฝังตำแหน่งที่จะใช้ ค่าเริ่มต้นคือ 1024dtype : JNP.DTYPE ประเภทข้อมูลที่จะใช้สำหรับการคำนวณ ค่าเริ่มต้นคือ jnp.float32causal : บูลีน ไม่ว่าจะใช้โหมดเชิงสาเหตุ (ไม่เร่งด่วน) หรือไม่ ค่าเริ่มต้นเป็น Truepolicy : สตริง นโยบายสำหรับการไล่ระดับสีด่าน ค่าเริ่มต้นคือ 'nothing_saveable'prevent_cse : บูลีน ไม่ว่าจะเป็นการป้องกันการกำจัด subexpression ทั่วไป (CSE) ค่าเริ่มต้นเป็น Falsefloat32_logits : บูลีน ไม่ว่าจะใช้ float32 สำหรับการคำนวณ logits ค่าเริ่มต้นเป็น Falseวิธีการ
วิธีการหลักของคลาส BlockwiseParallelJax คือวิธี forward ซึ่งดำเนินการส่งต่อของบล็อกหม้อแปลง
forward ( hidden_states , attention_mask , position_ids , deterministic = True , init_cache = False )hidden_states : jnp.ndarray เทนเซอร์อินพุตไปยังบล็อกหม้อแปลง มันควรจะมีรูปร่าง (batch_size, sequence_length, hidden_size)attention_mask : jnp.ndarray หน้ากากความสนใจสำหรับกลไกการดูแลตนเอง มันควรจะมีรูปร่าง (batch_size, 1, 1, sequence_length)position_ids : jnp.ndarray ID ตำแหน่งสำหรับการเข้ารหัสตำแหน่ง มันควรจะมีรูปร่าง (1, sequence_length)deterministic : บูลีน ไม่ว่าจะใช้โหมดกำหนด (ไม่มีการออกกลางคัน) หรือไม่ ค่าเริ่มต้นเป็น Trueinit_cache : บูลีน ไม่ว่าจะเริ่มต้นแคชสำหรับการถอดรหัสที่รวดเร็ว ค่าเริ่มต้นเป็น False วิธีนี้ส่งคืนเทนเซอร์เอาท์พุทของบล็อกหม้อแปลงซึ่งมีรูปร่างเหมือนกับ hidden_states
ตัวอย่างการใช้งาน
ตัวอย่างต่อไปนี้แสดงให้เห็นถึงวิธีการใช้คลาส BlockwiseParallelJax
# Initialize
from jax import random
import jax . numpy as jnp
from AttentionGrid import BlockwiseParallelJax
# Initialize transformer block
block = BlockwiseParallelJax (
q_chunk_size = 64 ,
k_chunk_size = 64 ,
hidden_size = 768 ,
num_heads = 12 ,
rotary_dim = 64 ,
intermediate_size = 3072 ,
)
# Create a batch of input tensors
key = random . PRNGKey ( 0 )
batch_size = 8
sequence_length = 128
hidden_states = random . normal ( key , ( batch_size , sequence_length , block . hidden_size ))
# Create attention mask
attention_mask = jnp . ones (( batch_size , 1 , 1 , sequence_length ))
# Create position ids
position_ids = jnp . arange ( sequence_length )[ None , :]
# Forward pass
output = block . forward ( hidden_states , attention_mask , position_ids )
print ( output . shape ) # prints: (8, 128, 768) FusedLandmarkAttention นี่คือคลาส Function pytorch ที่ห่อหุ้มฟังก์ชั่นไปข้างหน้าและย้อนกลับของกลไกความสนใจสถานที่สำคัญที่หลอมรวม
forward(ctx, q, k, v, n_prefix_q, sm_scale, block_size)ฟังก์ชั่นนี้ดำเนินการผ่านไปข้างหน้าของความสนใจสถานที่สำคัญที่หลอมรวม
ctx : วัตถุที่เราสามารถบันทึกตัวแปรสำหรับใช้ในผ่านไปข้างหลัง จัดทำโดยระบบอัตโนมัติของ Pytorchq : เคียวรีเทนเซอร์ สันนิษฐานว่าจะอยู่ติดกันและรูปร่างของมันควรจะเป็น (แบทช์, nheads, seqlen_q, d)k : คีย์เทนเซอร์ สันนิษฐานว่าจะอยู่ติดกันและรูปร่างของมันควรตรงกับรูปร่างของ Q เช่น (แบทช์, nheads, seqlen_k, d)v : ค่าเทนเซอร์ สันนิษฐานว่าจะอยู่ติดกันและรูปร่างของมันควรตรงกับรูปร่างของ Q และ K เช่น (แบทช์, nheads, seqlen_k, d)n_prefix_q : จำนวนคำนำหน้าในแบบสอบถามsm_scale : ปัจจัยการปรับขนาดที่ใช้ในการทำงานของ SoftMaxblock_size : ขนาดบล็อกสำหรับการดำเนินการบล็อกที่ชาญฉลาด o : เทนเซอร์เอาท์พุทจากผ่านไปข้างหน้าของกลไกความสนใจสถานที่สำคัญที่หลอมรวม backward(ctx, do)ฟังก์ชั่นนี้ดำเนินการผ่านไปข้างหลังของความสนใจที่เป็นแลนด์มาร์คที่ถูกหลอมรวมเช่นมันจะคำนวณการไล่ระดับสี
ctx : วัตถุที่เราสามารถดึงตัวแปรที่บันทึกไว้ในผ่านไปข้างหน้า จัดทำโดยระบบอัตโนมัติของ Pytorchdo : การไล่ระดับสีของการสูญเสียด้วยความเคารพต่อผลลัพธ์ของฟังก์ชันไปข้างหน้า Nonefused_landmark_attention(q, k, v, is_mem, sm_scale=None, block_size=64) ฟังก์ชั่นนี้เป็น wrapper ที่สะดวกสำหรับคลาส FusedLandmarkAttention
q : เคียวรีเทนเซอร์k : คีย์เทนเซอร์v : ค่าเทนเซอร์is_mem : เทนเซอร์บูลีนระบุว่าคู่คีย์ค่าแต่ละคู่ควรได้รับการปฏิบัติเป็นหน่วยความจำหรือไม่ มันควรมีความยาวเท่ากับความยาวลำดับของคีย์sm_scale : ปัจจัยการปรับขนาดที่ใช้ในการทำงานของ SoftMax หาก None มันจะถูกตั้งค่าเป็น 1.0 / sqrt(d)block_size : ขนาดบล็อกสำหรับการดำเนินการบล็อกที่ชาญฉลาด นี่คือตัวอย่างพื้นฐานของวิธีการใช้ฟังก์ชั่น fused_landmark_attention
import torch
from AttentionGrid import fused_landmark_attention
# Initialize some tensors
batch = 8
nheads = 12
seqlen = 128
d = 64
q = torch . randn ( batch , nheads , seqlen , d )
k = torch . randn ( batch , nheads , seqlen , d )
v = torch . randn ( batch , nheads , seqlen , d )
is_mem = torch . zeros ( seqlen , dtype = torch . bool )
# Call the function
output = fused_landmark_attention ( q , k , v , is_mem )
print ( output . shape ) # prints: (8, 12, 128, 64)ตัวอย่างนี้เป็นครั้งแรกที่เริ่มต้นเทนเซอร์บางอย่างเพื่อใช้เป็นคิวรีคีย์คีย์
และค่านิยม จากนั้นจะเรียกฟังก์ชั่น fused_landmark_attention และพิมพ์รูปร่างของเทนเซอร์เอาท์พุท
import torch
import torch . nn as nn
from AttentionGrid import DilatedAttention
# Replace this with your correct GPU device
device = "cuda:0"
dtype = torch . float16
# Create an instance of DilatedAttention
d_model = 512
num_heads = 8
dilation_rate = 2
segment_size = 64
dropout = 0.2 # Specify the dropout rate
attention = DilatedAttention (
d_model = d_model ,
num_heads = num_heads ,
dilation_rate = dilation_rate ,
segment_size = segment_size ,
dropout = dropout ,
). to ( device , dtype = dtype )
# Create some dummy input data
batch_size = 16
seq_len = 128
input_dim = d_model
inputs = torch . randn ( batch_size , seq_len , input_dim , device = device , dtype = dtype )
# Forward pass
outputs = attention ( inputs )
# Print the output shape
print ( outputs . shape ) # Expected: [batch_size, seq_len, d_model] ในตัวอย่างด้านบนเราสร้างตัวอย่างของคลาส DilatedAttention ด้วย hyperparameters ที่ระบุ จากนั้นเราจะสร้างข้อมูลอินพุตจำลองบางอย่างและส่งผ่านกลไกความสนใจเพื่อให้ได้ผลลัพธ์ ในที่สุดเราพิมพ์รูปร่างของเทนเซอร์เอาท์พุท
คลาส DilatedAttention ใช้ความสนใจที่ขยายตัวซึ่งขยายสนามที่เอาใจใส่อย่างทวีคูณเมื่อระยะห่างระหว่างโทเค็นเติบโต มันสืบทอดมาจาก torch.nn.Module และสามารถใช้แทนกลไกความสนใจมาตรฐานในแบบจำลองหม้อแปลง
d_model (int): มิติของการฝังอินพุตและเอาต์พุตnum_heads (int): จำนวนหัวความสนใจdilation_rate (int): อัตราการขยายสำหรับการขยายลำดับอินพุตsegment_size (int): ขนาดของแต่ละเซ็กเมนต์หลังจาก sparsificationdropout (ลอย, เป็นตัวเลือก): ความน่าจะเป็นแบบออกกลางคันที่จะนำไปใช้กับเอาต์พุตความสนใจ ค่าเริ่มต้น: 0.0 (ไม่มีการออกกลางคัน)x (เทนเซอร์): เทนเซอร์อินพุตของรูปร่าง (batch_size, seq_len, d_model)output (เทนเซอร์): เทนเซอร์เอาท์พุทของรูปร่าง (batch_size, seq_len, d_model) โปรดทราบว่าเทนเซอร์อินพุตควรอยู่ในอุปกรณ์ที่ถูกต้อง (เช่น GPU) และมีประเภทข้อมูลที่เหมาะสม ( dtype )
