NVIDIA เปิดตัว BioNeMo Inference Runtime เร่งความเร็ว

NVIDIA ประกาศเปิดตัว BioNeMo Inference Runtime หรือ BioIR อย่างเป็นทางการ โดยเป็นรันไทม์ที่ออกแบบมาเพื่อเร่งความเร็วโมเดลโครงสร้างชีววิทยาเชิงคำนวณและการพับโปรตีนโดยเฉพาะ พร้อมเปิดให้ใช้งานแล้วผ่านทาง GitHub repository พร้อม wheel ที่บรรจุ precompiled CUBINs ไว้เรียบร้อยแล้ว ความต้องการของระบบในการใช้งานรันไทม์นี้จำเป็นต้องใช้ Python 3.12 ขึ้นไป, การ์ดจอ NVIDIA และไดรเวอร์ที่รองรับ, โมเดลเช็คพอยต์ที่จัดเตรียมไว้ รวมถึง A3M MSAs ต่อเชน โดยไม่ต้องพึ่งพา nvcc, CUDA source, CMake หรือ CUDA toolkit แต่อย่างใด

BioIR มุ่งเน้นการจัดการกับการปฏิบัติการที่สแต็กอินเฟอเรนซ์แบบทั่วไปยังไม่สามารถปรับแต่งให้เหมาะสมได้อย่างเต็มที่ ซึ่งรวมถึงสเกลของ Pairformer และ Evoformer, การดำเนินการแบบสามเหลี่ยม (triangle operations), pairwise attention, diffusion transformers และโมดูลระดับอะตอม โดยตัวโมเดลจะยังคงสถานะเป็นวัตถุ torch.nn.Module ปกติ ทำให้ไม่มีขั้นตอนการสร้างเอนจิน, ขั้นตอนการส่งออก (export step) หรืออาร์ติแฟกต์แยกต่างหากระหว่างเช็คพอยต์กับการส่งผ่านข้อมูลไปข้างหน้า (forward pass)

ภาพประกอบจากคลังภาพสต็อก ไม่ใช่ภาพจากเหตุการณ์จริง

สำหรับการใช้งาน นักพัฒนาสามารถเลือกใช้งานได้ 2 วิธีหลัก ได้แก่ เอนด์ทูเอนด์โปรเซสเซอร์ (End-to-end processor) ซึ่งจะเคลื่อนย้าย InputRequest ผ่านขั้นตอนการแยกวิเคราะห์ การทำโทเค็นไนเซชัน การสร้างฟีเจอร์ การอนุมานด้วย GPU และการเขียนไฟล์ PDB หรือ mmCIF ส่วนวิธีที่สองคือการบูรณาการ PyTorch โดยตรง (Direct PyTorch integration) ซึ่งให้นักพัฒนาสามารถสร้างโมเดลที่รองรับหรือนำโมดูลที่ปรับแต่งแล้วกลับมาใช้ใหม่ภายในโค้ดของตนเองได้ ตัวบทเรียนสอนใช้งาน (tutorial) ได้แสดงเส้นทางโปรเซสเซอร์ด้วย Boltz-2 โดยแต่ละเชนโปรตีนจะต้องการ A3M MSA และสามารถรองรับ MSA ทั้งแบบจับคู่หรือไม่จับคู่สำหรับอินพุตที่มีเชนโปรตีนไม่เหมือนกันหลายเชน

ในแง่ของการประเมินประสิทธิภาพระดับโมเดลฟอร์เวิร์ด รายงานการเบนช์มาร์กช่วงแรกของ NVIDIA แสดงให้เห็นความเร็วเฉลี่ยเชิงเรขาคณิตเมื่อเทียบกับฐานข้อมูล OSS torch.compile โดยบนชิป H100 ทำความเร็วได้ 1.55x สำหรับ OpenFold3, 1.78x สำหรับ Boltz2 และ 2.56x สำหรับ OpenFold2 monomer ขณะที่ตัวเลขบนชิป H200 อยู่ที่ระดับใกล้เคียงกันคือ 1.54x, 1.75x และ 2.61x ซึ่งวัดจากอินพุต 17 รายการที่มีความยาวตั้งแต่ 29 ถึง 1,734 เรสิดิว

เพื่อให้เห็นภาพการส่งมอบงานแบบเอนด์ทูเอนด์ ทีมงาน NVIDIA ได้ทำการทดสอบเปรียบเทียบกับเป้าหมายไดเมอร์ของมนุษย์จำนวน 1,000 รายการ ที่มีความยาวลำดับสังกะตังรวมกันต่ำกว่า 2,800 เรสิดิว การทดสอบนี้นำ Boltz-2 ที่เร่งความเร็วด้วย BioIR ไปเทียบกับ Boltz-2 แบบโอเพนซอร์สที่คอมไพล์ด้วย torch บนจีพียู 8xH100 80GB โดยใช้เป้าหมาย MSA ที่จัดเตรียมไว้ สูตรการอนุมาน (3 recycles, 200 sampling steps, 5 diffusion samples) และการกำหนดค่า GPU ที่เหมือนกันทุกประการ

"These numbers are folding-stage measurements specific to this dataset and hardware. They exclude MSA generation, preprocessing CPU allocations, storage, data transfer, and retries."

NVIDIA Research Blog

ภาพประกอบจากคลังภาพสต็อก ไม่ใช่ภาพจากเหตุการณ์จริง

เมื่อทำการประมาณการขยายผลการเบนช์มาร์คแบบเชิงเส้นไปยังเป้าหมายที่เทียบเคียงกันได้ถึง 1 ล้านรายการ ประเมินว่า BioIR จะใช้พลังงานราว 11 เมกะวัตต์-ชั่วโมง (MWh) เทียบกับ 35 เมกะวัตต์-ชั่วโมงของตัวเลือกแบบสาธารณะเมื่อคำนวณตามสัดส่วน TDP ของ 8-GPU และหากใช้ค่าเทียบเท่าพลังงานสูงสุดของฟูลโหนด ตัวเลขจะอยู่ที่ 21 เมกะวัตต์-ชั่วโมงเทียบกับ 64 เมกะวัตต์-ชั่วโมง ซึ่งตัวเลขเหล่านี้เป็นเพียงการประเมินกำลังไฟเฉพาะช่วงการพับโปรตีนสำหรับอุปกรณ์ไอทีเท่านั้น ไม่รวมค่าใช้จ่ายด้านโครงสร้างพื้นฐานศูนย์ข้อมูล เช่น PUE แต่การประหยัดพลังงานระดับ 23 ถึง 43 เมกะวัตต์-ชั่วโมงต่อการประมวลผล 1 ล้านเป้าหมาย ถือเป็นตัวเลขที่มีความสำคัญอย่างยิ่งสำหรับการทำแคมเปญระดับโปรโตม

ที่มา: MarkTechPost