สร้างเวิร์กโฟลว์ AI แบบสหพันธ์ (Federated AI) ที่รองรับหลายรูปแบบข้อมูล (Multimodal) ด้วย NVIDIA FLARE
ในยุคปัจจุบัน โมเดล AI ที่สามารถประมวลผลได้ทั้งภาพและข้อความ (Vision-Language Models - VLMs) มีความสามารถหลากหลาย ตั้งแต่การตอบคำถามจากภาพ การสร้างคำบรรยายภาพ ไปจนถึงการให้เหตุผลระหว่างรูปภาพและข้อความ อย่างไรก็ตาม การจะปรับปรุงโมเดลเหล่านี้ให้มีประสิทธิภาพสูงขึ้น มักต้องการข้อมูลจำนวนมากที่อาจกระจายตัวอยู่ตามสถาบันหรือองค์กรต่างๆ ซึ่งทำให้การรวบรวมข้อมูลทั้งหมดมาไว้ที่ส่วนกลางเป็นเรื่องท้าทาย
การเรียนรู้แบบสหพันธ์ (Federated Learning) เป็นแนวทางที่ช่วยให้สามารถฝึกฝนโมเดล AI ร่วมกันได้ โดยไม่ต้องย้ายข้อมูลดิบออกจากแหล่งกำเนิด แต่สำหรับ VLM นั้น ความท้าทายไม่ได้อยู่ที่การจัดการฝึกฝนเท่านั้น แต่ยังรวมถึงการที่แต่ละไซต์ (Client) อาจมีส่วนผสมของข้อมูลและงาน (Task) ที่แตกต่างกัน และขนาดของการอัปเดตโมเดลก็อาจใหญ่เกินกว่าที่แบนด์วิดท์เครือข่ายหรือหน่วยความจำของเซิร์ฟเวอร์จะรองรับได้
บทความนี้จะเน้นไปที่การตัดสินใจสำคัญสองประการในการออกแบบเวิร์กโฟลว์ AI แบบสหพันธ์ที่รองรับหลายรูปแบบข้อมูล:
- ควรส่งข้อมูลสถานะโมเดลส่วนใดข้ามเครือข่ายบ้าง?
- จะถ่ายโอนและรวมการอัปเดตโมเดลเหล่านั้นอย่างมีประสิทธิภาพได้อย่างไร?
เราจะแสดงให้เห็นว่า NVIDIA FLARE สามารถช่วยประสานงานการฝึกฝนแบบสหพันธ์ได้อย่างไร และจัดการกับการอัปเดตโมเดลขนาดใหญ่ผ่านเทคนิคต่างๆ เช่น การทำให้วัตถุขนาดใหญ่เป็นภายนอก (Externalization), การสตรีมเทนเซอร์ (Tensor Streaming) และการรวมข้อมูลโดยใช้ดิสก์เป็นส่วนรองรับ (Disk-backed Aggregation)
ความท้าทายในการทำ Federated VLM
ในการทดลอง Vision-Language แบบรวมศูนย์ (Centralized) ข้อมูลต่างๆ เช่น รูปภาพ คำบรรยาย ภาพตัวอย่างคำถาม-คำตอบ หรือพรอมต์สำหรับสร้างข้อความ สามารถป้อนเข้าสู่ไปป์ไลน์การฝึกฝนเดียวได้ แต่ในสภาพแวดล้อมแบบสหพันธ์ ข้อมูลเหล่านี้จะกระจายตัวอยู่ตามไซต์ต่างๆ ซึ่งแต่ละไซต์อาจมีข้อมูล งาน หรือข้อจำกัดในการทำงานที่แตกต่างกัน
ปัญหานี้ก่อให้เกิดความท้าทายทางวิศวกรรมสองประการ:
- ความหลากหลายของข้อมูลและงาน: แต่ละไซต์อาจฝึกฝนบนชุดงานหรือรูปแบบข้อมูลที่แตกต่างกัน ดังนั้น เวิร์กโฟลว์ต้องกำหนดอย่างชัดเจนว่าแต่ละไคลเอนต์ควรจะอัปเดตส่วนใดของโมเดล และจะรวมการอัปเดตเหล่านั้นเข้าด้วยกันอย่างไร
- ขนาดของการอัปเดตโมเดล: การส่งพารามิเตอร์โมเดลทั้งหมดอาจใช้ทรัพยากรในการแปลงเป็นซีเรียล (Serialization) ถ่ายโอน และเก็บไว้ในหน่วยความจำเซิร์ฟเวอร์เป็นจำนวนมาก
ดังนั้น การตัดสินใจแรกที่สำคัญคือ "จะทำ Federation กับอะไร?" บางแนวทางอาจเลือกแลกเปลี่ยนความรู้ที่กลั่นกรองมา (Distilled Knowledge) แทนที่จะเป็นน้ำหนักโมเดล (Model Weights) ในขณะที่บางแนวทางจะตรึง (Freeze) ส่วน backbone ของโมเดลที่ได้รับการฝึกฝนล่วงหน้าไว้ แล้วจึงรวมเฉพาะส่วนประกอบที่สามารถฝึกฝนได้ซึ่งมีขนาดเล็กกว่า
เมื่อจำเป็นต้องมีการอัปเดตที่ใหญ่ขึ้น ระบบก็ต้องรองรับการสตรีมและการรวมข้อมูลที่ใช้หน่วยความจำอย่างมีประสิทธิภาพ NVIDIA FLARE ซึ่งเป็น SDK และเฟรมเวิร์กแบบโอเพนซอร์สที่ขยายได้สำหรับ Federated Learning และการประมวลผลร่วมกัน สามารถรองรับรูปแบบการสื่อสารทั้งแบบ Parameter-Efficient (เช่น การใช้ Adapter) และแบบ Full-Model ได้
การประสานงานการฝึกฝนระหว่างไคลเอนต์
NVIDIA FLARE แยกการประสานงานส่วนกลาง (Global Coordination) ออกจากการดำเนินการในแต่ละไซต์ (Local Execution) อย่างชัดเจน โดยเซิร์ฟเวอร์จะทำหน้าที่กำหนดตารางรอบการฝึกฝน (Rounds) และรวบรวมการอัปเดต ในขณะที่แต่ละไคลเอนต์จะทำการฝึกฝนหรือประเมินผลโดยใช้ข้อมูลในเครื่องของตนเอง การประมวลผลล่วงหน้า (Preprocessing) การสร้างพรอมต์ (Prompt Construction) และการจัดกลุ่มข้อมูล (Batching) จะยังคงอยู่ภายในไคลเอนต์
API ที่ชื่อว่า Recipe API ของ NVIDIA FLARE เป็นจุดเริ่มต้นที่กระชับ โดยสูตร FedAvg จะจับคู่โมเดลเข้ากับสคริปต์การฝึกฝนของไคลเอนต์ ซึ่งสูตรเดียวกันนี้สามารถนำไปใช้ในการจำลอง (Simulation) หรือการติดตั้งใช้งานจริงแบบหลายไซต์ได้
ก่อนที่จะเริ่มการเขียนโค้ดโมเดล ควรมีการกำหนด "สัญญาการอัปเดตของไคลเอนต์" (Client Update Contract) ซึ่งระบุว่า:
- ข้อมูลส่วนใดจะยังคงอยู่ภายในเครื่อง (Local)
- ข้อมูลส่วนใดที่อาจถูกส่งออกจากไซต์ได้
- ส่วนประกอบโมเดลใดบ้างที่แต่ละไคลเอนต์สามารถอัปเดตได้
- เมตริกใดบ้างที่จะส่งกลับไปยังเซิร์ฟเวอร์
- เมื่อไคลเอนต์อัปเดตส่วนประกอบโมเดลที่แตกต่างกัน สัญญาดังกล่าวควรกำหนดวิธีการรวมการอัปเดตระดับส่วนประกอบเหล่านั้นเข้าด้วยกันด้วย
การเคลื่อนย้ายและรวมการอัปเดตโมเดลขนาดใหญ่ให้มีประสิทธิภาพ
วิธีการพื้นฐานที่นิยมใช้ในการฝึกฝน VLM แบบสหพันธ์ คือการปรับแต่ง (Fine-tuning) และรวมพารามิเตอร์โมเดลทั้งหมด ซึ่งจะส่งผลให้การอัปเดตโมเดลมีขนาดใหญ่ เมื่อมีไคลเอนต์จำนวนมากเข้าร่วมในการทำ Federation จะเกิดแรงกดดันต่อหน่วยความจำสองประการหลักๆ คือ การแปลงเป็นซีเรียลและถ่ายโอนการอัปเดตหนึ่งครั้ง และการเก็บการอัปเดตจากหลายไคลเอนต์ไว้ในหน่วยความจำระหว่างการรวม NVIDIA FLARE รองรับคุณสมบัติหลายประการเพื่อจัดการกับความท้าทายนี้
1. ทำให้วัตถุขนาดใหญ่เป็นภายนอก (Externalize Large Objects)
NVIDIA FLARE สามารถแทนที่วัตถุขนาดใหญ่ในข้อความด้วยการอ้างอิงที่มีน้ำหนักเบา และถ่ายโอนข้อมูลจริงแยกต่างหาก วิธีนี้ช่วยให้ข้อความควบคุม (Control Message) มีขนาดเล็ก และรองรับเพย์โหลด (Payload) ที่เกินขีดจำกัดของข้อความที่ถูกแปลงเป็นซีเรียลตามปกติ โมดูลที่สร้างไว้ล่วงหน้า (Built-in decomposers) รองรับ PyTorch tensors, NumPy arrays และโครงสร้าง FLARE ทั่วไป การสร้าง decomposer แบบกำหนดเองจะจำเป็นเฉพาะสำหรับประเภทวัตถุที่เฉพาะเจาะจงของแอปพลิเคชันเท่านั้น
สำหรับเวิร์กโฟลว์ PyTorch, FLARE Tensor Downloader จะสตรีมเทนเซอร์แบบค่อยเป็นค่อยไปโดยใช้โปรโตคอลแบบดึง (Pull-based protocol) โดยจะมีการแปลงเป็นซีเรียลเพียงบางส่วนในแต่ละครั้ง ซึ่งช่วยลดหน่วยความจำสูงสุด (Peak Memory) ระหว่างการกระจายโมเดล ขนาดของชิ้นส่วน (Chunk size) สามารถปรับได้เพื่อสร้างสมดุลระหว่างค่าใช้จ่ายในการร้องขอ (Request Overhead) กับหน่วยความจำต่อชิ้นส่วน
2. การใช้ดิสก์ช่วยในการรวมข้อมูล (Offload Aggregation to Disk)
การสตรีมช่วยลดแรงกดดันต่อหน่วยความจำระหว่างการถ่ายโอน แต่เซิร์ฟเวอร์อาจยังคงต้องเก็บการอัปเดตจากหลายไคลเอนต์ไว้ในหน่วยความจำระหว่างการรวม ซึ่งอาจทำให้หน่วยความจำสูงสุดของเซิร์ฟเวอร์เพิ่มขึ้นตามจำนวนไคลเอนต์ ใน NVIDIA FLARE เวอร์ชัน 2.8.0 เป็นต้นมา Tensor Disk Offload จะเขียนการอัปเดต FedAvg ของ PyTorch ที่เข้ามาไปยังไฟล์ safetensors ชั่วคราว และโหลดข้อมูลเมื่อจำเป็น วิธีนี้ช่วยป้องกันไม่ให้หน่วยความจำ CPU เพิ่มขึ้นตามสัดส่วนเชิงเส้น
กลไกเหล่านี้จะช่วยเสริมการลดเพย์โหลดจากการฝึกฝนแบบใช้ Adapter (Adapter-based training) เพื่อเปิดใช้งานการฝึกฝนแบบ Full-model, การใช้ Adapter ที่มีขนาดใหญ่ขึ้น หรือ Federated Learning ที่มีไคลเอนต์จำนวนมาก
การทำ Federation สำหรับ Adapter ขนาดเล็กบน VLM ที่ตรึงไว้ด้วย FedUMM
FedUMM เป็นตัวอย่างที่จับต้องได้ในการลดขนาดข้อมูลที่ต้องส่งข้ามเครือข่ายใน NVIDIA FLARE โดยแต่ละไคลเอนต์ที่จำลองขึ้นจะรักษา backbone ของโมเดล BLIP ที่ถูกตรึงไว้ และทำการฝึกฝน Adapter ในระดับท้องถิ่น NVIDIA FLARE จะทำหน้าที่ประสานงานรอบการฝึกฝนและรวมเฉพาะการอัปเดต Adapter เท่านั้น FedUMM ได้รับการออกแบบมาเพื่อความยืดหยุ่น โดยมีตัวเข้ารหัสเฉพาะสำหรับแต่ละรูปแบบข้อมูล (Modality-specific encoders) สำหรับภาพ เสียง และข้อความ แม้ว่าการทดลองในปัจจุบันจะเน้นไปที่การประมวลผลภาพและข้อความก็ตาม
การทดลองที่รายงานได้ประเมินโมเดลบนชุดข้อมูล VQA v2 และ GenEval ภายใต้สภาวะความไม่สม่ำเสมอของข้อมูลที่ควบคุมโดย Dirichlet (Dirichlet-controlled heterogeneity) ด้วยไคลเอนต์สูงสุด 16 ราย ในการเปรียบเทียบกับไคลเอนต์ 8 ราย การทำ Federation เฉพาะ Adapter ช่วยลดการสื่อสารต่อไคลเอนต์จาก 28.6 GB เหลือเพียง 0.094 GB ต่อรอบ และปรับปรุงประสิทธิภาพบน VQA v2 ได้ 0.7 คะแนน เมื่อเทียบกับ FedAvg แบบ Full-model ที่ไคลเอนต์ 8 ราย ประสิทธิภาพยังคงอยู่ที่ประมาณ 97% ของการอ้างอิงแบบรวมศูนย์ในทั้งสองเกณฑ์มาตรฐาน
การประเมินนี้ใช้ไซต์ที่จำลองขึ้น พาร์ติชันสังเคราะห์ และเกณฑ์มาตรฐานสาธารณะทั่วไป ซึ่งไม่ได้เป็นการพิสูจน์ประสิทธิภาพทางคลินิกหรือการรับประกันความเป็นส่วนตัวอย่างเป็นทางการ แต่เป็นการแสดงให้เห็นว่าข้อมูลดิบยังคงอยู่ในระดับท้องถิ่นภายในเวิร์กโฟลว์แบบสหพันธ์ที่จำลองขึ้น
FedUMM ช่วยลดภาระของระบบที่ต้นทางโดยการแลกเปลี่ยนเฉพาะ Adapter ขนาดเล็กเท่านั้น อย่างไรก็ตาม ไม่ใช่ทุกเวิร์กโฟลว์ AI ที่จะสามารถทำเช่นนี้ได้ เมื่อไคลเอนต์ต้องส่งการอัปเดตที่ใหญ่ขึ้น การสตรีมเทนเซอร์จะช่วยลดแรงกดดันต่อหน่วยความจำระหว่างการถ่ายโอน และเมื่อเซิร์ฟเวอร์ต้องรวมการอัปเดตจากไคลเอนต์จำนวนมาก การใช้ดิสก์ช่วยในการรวมข้อมูลจะช่วยลดปริมาณข้อมูลที่ต้องเก็บไว้ในหน่วยความจำ
Checklist สำหรับการออกแบบ Federated Multimodal AI Workflows
เมื่อออกแบบเวิร์กโฟลว์ AI แบบสหพันธ์ที่รองรับหลายรูปแบบข้อมูล ควรพิจารณาทั้งสิ่งที่แต่ละไคลเอนต์ควรมีส่วนร่วม และวิธีการที่การอัปเดตจะเคลื่อนที่ผ่านระบบ Checklist ต่อไปนี้สรุปการตัดสินใจที่สำคัญเพื่อสร้างสมดุลระหว่างคุณภาพของโมเดล ต้นทุนการสื่อสาร และข้อกำหนดด้านหน่วยความจำของระบบ:
- กำหนดสัญญาการอัปเดต (Define the Update Contract): ตัดสินใจว่าอะไรจะยังคงอยู่ภายในเครื่อง (Local) อะไรที่แต่ละไคลเอนต์จะส่งกลับมา และจะรวมการอัปเดตเหล่านั้นเข้าด้วยกันอย่างไร
- ลดขนาดเพย์โหลด (Minimize the Payload): เลือกแลกเปลี่ยน Adapter ที่มีน้ำหนักเบาเมื่อเป็นไปได้ และใช้การอัปเดตแบบ Full-model เฉพาะเมื่อจำเป็นสำหรับงานนั้นๆ
- เลือกวิธีการเคลื่อนย้ายและรวมการอัปเดต (Choose how Updates Move and Aggregate): ใช้ Externalization และ Tensor Streaming สำหรับการอัปเดตขนาดใหญ่ที่อยู่ในหน่วยความจำ ร่วมกับการใช้ Disk Offload บนเซิร์ฟเวอร์ เมื่อการรวมการอัปเดตหลายรายการอาจเกินขีดจำกัดหน่วยความจำของเซิร์ฟเวอร์
- ประเมินแบบ End-to-End (Evaluate End-to-End): วัดคุณภาพของโมเดล ควบคู่ไปกับการสื่อสาร เวลาประมวลผล การใช้หน่วยความจำ ความไม่สม่ำเสมอของข้อมูล และข้อผิดพลาดที่อาจเกิดขึ้น
เริ่มต้นสร้าง Federated Multimodal AI Workflows
เริ่มต้นด้วยการกำหนดสัญญาการอัปเดตสำหรับเวิร์กโฟลว์ของคุณ: อะไรจะยังคงอยู่ภายในเครื่อง, สถานะโมเดลใดบ้างที่แต่ละไคลเอนต์สามารถส่งกลับมาได้, และไคลเอนต์ใดบ้างที่ควรมีส่วนร่วมในการรวมข้อมูลแต่ละครั้ง ใช้ NVIDIA FLARE Recipe API เพื่อนำเวิร์กโฟลว์ไปใช้งานและตรวจสอบความถูกต้องในการจำลองด้วยจำนวนไคลเอนต์ที่คาดหวัง
ถัดไป เลือกกลไกการจัดการเพย์โหลดที่ตรงกับคอขวด (Bottleneck) ของคุณ ใช้ Large-Object Externalization และ FLARE Tensor Downloader สำหรับการอัปเดตโมเดลขนาดใหญ่ ควบคู่ไปกับการใช้ Tensor Disk Offload เมื่อหน่วยความจำในการรวมข้อมูลฝั่งเซิร์ฟเวอร์กลายเป็นข้อจำกัด
สำหรับตัวอย่างที่จับต้องได้ของการฝึกฝนแบบใช้
ขอบคุณ แหล่งข้อมูล
https://developer.nvidia.com/blog/building-federated-multimodal-ai-workflows-with-nvidia-flare/