২ সেকেন্ডের কম সময়ে RL পোস্ট-ট্রেনিংয়ের জন্য ওয়েট ট্রান্সফার

আমরা সম্প্রতি কিমি-কে২ (১টি প্যারামিটার) এর জন্য ১.৩-সেকেন্ডে ক্রস-মেশিন প্যারামিটার আপডেট অর্জন করেছি, যা ২৫৬টি ট্রেনিং জিপিইউ (BF16) থেকে ১২৮টি ইনফারেন্স জিপিইউ (FP8)-তে ওয়েট ট্রান্সফার করে।

অ্যাসিনক্রোনাস রিইনফোর্সমেন্ট লার্নিং ফাইন-টিউনিংয়ে, ট্রেনিং এবং ইনফারেন্স আলাদা জিপিইউতে চলে। প্রতিটি ট্রেনিং ধাপের পর, নতুন ওয়েট ইনফারেন্স নোডগুলোতে পাঠাতে হয়। ট্রিলিয়ন-প্যারামিটার মডেলের ক্ষেত্রে অনেক বিদ্যমান ফ্রেমওয়ার্কে কয়েক সেকেন্ড—বা মিনিটও—লেগে যায়।

RDMA পয়েন্ট-টু-পয়েন্ট কমিউনিকেশন-এর সুবিধা নিয়ে, আমরা ইনফারেন্স ইঞ্জিন পরিবর্তন না করেই এবং কোড লেখা ও রক্ষণাবেক্ষণ সহজ করে, ওয়েট ট্রান্সফারকে অত্যন্ত দ্রুত করে তুলতে সক্ষম হয়েছি।

আমাদের সমাধানটি RDMA WRITE-এর ওপর ভিত্তি করে তৈরি, যা একটি একমুখী প্রিমিটিভ যেখানে সোর্স সরাসরি ডেস্টিনেশনের জিপিইউ মেমোরিতে লেখে।

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

ডেস্টিনেশন সাইড ট্রান্সফারের জন্য কোনো নোটিফিকেশনও পাবে না। এটি আমাদের ট্রেনিং নোড দ্বারা চালিত কম লেটেন্সি, উচ্চ থ্রুপুট, জিরো-কপি ট্রান্সফার প্রদান করে, যেখানে ইনফারেন্স নোডগুলোতে কোনো কন্ট্রোল লজিক থাকে না।

মেটাডেটা সংগ্রহ – কন্ট্রোলার সমস্ত ট্রেনিং এবং ইনফারেন্স জিপিইউ থেকে প্যারামিটার মেটাডেটা সংগ্রহ করে।

শিডিউল গণনা – কন্ট্রোলার একটি স্ট্যাটিক ওয়েট ট্রান্সফার শিডিউল গণনা করে, যা ম্যাপ করে কোন ট্রেনিং জিপিইউ কোন প্যারামিটারটি কোন ইনফারেন্স জিপিইউতে পাঠাবে এবং কোন ক্রমে পাঠাবে।

শিডিউল বিতরণ – কন্ট্রোলার সমস্ত ট্রেনিং জিপিইউতে শিডিউল পাঠায়।

এক্সিকিউশন – প্রতিটি ট্রেনিং ধাপের পরে, কন্ট্রোলার ট্রেনিং জিপিইউগুলোকে ট্রান্সফার শুরু করার জন্য সংকেত দেয়।

হাই-লেভেল ওয়ার্কফ্লো সংজ্ঞায়িত হওয়ার পর, মূল চ্যালেঞ্জ হলো ট্রিলিয়ন-প্যারামিটার স্কেলে কীভাবে দক্ষতার সাথে ওয়েট ট্রান্সফার এক্সিকিউট করা যায়। এখানে আমরা এক্সিকিউশন পাথের বিবরণ দিচ্ছি।

ট্রেনিংয়ে প্যারামিটারগুলো FSDP প্লেসমেন্ট অনুযায়ী বিতরণ করা হয়। full_tensor() ব্যবহার করে, একটি DeviceMesh-এর সমস্ত জিপিইউ সম্পূর্ণ প্যারামিটার পুনর্গঠন করতে পারে, তাই সবাই ওয়েট ট্রান্সফারের জন্য সোর্স হিসেবে কাজ করতে পারে।

একাধিক ডিসজয়েন্ট ডিভাইসমেশ একটি মেশ গ্রুপ গঠন করে। যেহেতু একই গ্রুপের ডিভাইসমেশগুলো ডিসজয়েন্ট, তাই তাদের ট্রান্সফারগুলো একে অপরের সাথে সাংঘর্ষিক হয় না এবং সম্পূর্ণ সমান্তরালে চলতে পারে। মেশ গ্রুপগুলোর মধ্যে, আমরা ক্রম বজায় রাখতে একটি গ্লোবাল ব্যারিয়ার যুক্ত করি।

ওজন-স্থানান্তর পাইপলাইন হোস্ট-টু-ডিভাইস কপি, টেনসর অ্যাসেম্বলি, ফিউশন ও কোয়ান্টাইজেশন, এবং MeshGroups জুড়ে RDMA ট্রান্সফারের সঙ্গে ওভারল্যাপ করে.

আমরা প্রতিটি প্যারামিটার টেন্সরের ট্রান্সফারকে একটি টাস্ক হিসেবে বিবেচনা করি। ওয়েট ট্রান্সফার প্রক্রিয়ায় বিভিন্ন ধরনের হার্ডওয়্যার সোর্স ব্যবহার করা হয়, তাই আমরা একটি ওয়েট ট্রান্সফার টাস্কে বিভিন্ন পাইপলাইন স্টেজে বিভক্ত করি যা সময়ে ওভারল্যাপ করে:

হোস্ট থেকে ডিভাইস মেমসিপি — যদি FSDP সিপুতে ওয়েট অফলোড করে

প্যারামিটার প্রস্তুতি — full_tensor() দিয়ে সম্পূর্ণ ওয়েট পুনর্গঠন করুন, প্রজেকশন ফিউশন প্রয়োগ করুন, প্রয়োজন হলে কোয়ান্টাইজ করুন।

RDMA ট্রান্সফার — রিমোট ইনফারেন্স জিপিইউ মেমোরিতে জিরো-কপি রাইট

গ্লোবাল ব্যারিয়ার — সমস্ত full_tensor() কল সম্পন্ন হওয়ার পর, ইথারনেটের মাধ্যমে GLOO ব্যবহার করে মেশ গ্রুপ জুড়ে সিঙ্ক্রোনাইজ করুন।

বাস্তবায়নে, আমরা প্রতিটি পাইপলাইন স্টেজের জন্য টাস্কের একটি FIFO সারি বজায় রাখি। সারির মাথার টাস্কটি যখন স্টেজ সম্পন্ন করে, তখন এটি পরবর্তী স্টেজ সারির টেলে স্থানান্তরিত হয়।

full_tensor() এবং অন্যান্য জিপিইউ অপারেশন অতিরিক্ত জিপিইউ মেমোরি ব্যবহার করে। মেমোরি শেষ হয়ে যাওয়ার ত্রুটি এড়াতে, আমরা একটি টাস্কের এক্সিকিউশন তখনই শুরু করি যদি বর্তমান অন-দ্য-ফ্লাই টাস্কগুলো কনফিগারযোগ্য ওয়াটারমার্কের চেয়ে কম সাময়িক জিপিইউ মেমোরি দখল করে।

কিছু ডিজাইন পছন্দের কারণে আমাদের সিস্টেমটি সাধারণ ওপেন-সোর্স সমাধানের তুলনায় চালানো উল্লেখযোগ্যভাবে দ্রুত এবং রক্ষণাবেক্ষণ করা সহজ।

Rank0-ভিত্তিক ওজন স্থানান্তর ইনফারেন্স নোডে বিতরণ করার আগে ওজনগুলো সংগ্রহ করে.

পয়েন্ট-টু-পয়েন্ট যোগাযোগ উভয় ট্রেনিং নোড থেকে ওজনগুলো ইনফারেন্স নোডের নেটওয়ার্ক ইন্টারফেসে স্থানান্তর করে.

একটি সাধারণ প্যাটার্ন হলো সমস্ত প্যারামিটারকে র‍্যাঙ্ক-০ জিপিইউ-এর মাধ্যমে পাঠানো: ট্রেনিং র‍্যাঙ্ক-০ তে গ্যাদার করা, ইনফারেন্স র‍্যাঙ্ক-০ তে পাঠানো, তারপর আবার স্ক্যাটার করা। এটি দ্রুত একটি বাধা হয়ে দাঁড়ায়, যা একটি একক জিপিইউ-এর PCIe ব্যান্ডউইথ এবং NIC (যেমন, ৪০০ Gbps ≈ ৫০ GB/s) দ্বারা সীমিত।

এর বিপরীতে, আমাদের পয়েন্ট-টু-পয়েন্ট সেটআপ প্রতিটি ট্রেনিং জিপিইউকে সরাসরি প্রতিটি ইনফারেন্স জিপিইউতে পাঠাতে দেয়, যা একটি লিঙ্কের পরিবর্তে সম্পূর্ণ নেটওয়ার্ক ফ্যাব্রিককে স্যাচুরেট করে।

RPC-ভিত্তিক স্থানান্তরে একটি ইনফারেন্স-সাইড বাফার ও worker কল ব্যবহার করা হয়, আর one-sided স্থানান্তর ইনফারেন্স-worker অংশগ্রহণ ছাড়াই সরাসরি ওজন লিখে দেয়.

কিছু সিস্টেম প্রতিটি টেনসরের জন্য inference engine-এর update_weight() মেথডে কল করার ওপর নির্ভর করে। এর মানে হলো inference কোডে আক্রমণাত্মক পরিবর্তন, পাশাপাশি RPC, serialization, এবং control-plane coordination থেকে অতিরিক্ত overhead।

RDMA WRITE প্রিমিটিভের মাধ্যমে, আমরা অতিরিক্ত কপি ছাড়াই ইনফারেন্স জিপিইউ মেমোরিতে নীরবে ওয়েট আপডেট করি। কোনো কন্ট্রোল প্লেন মেসেজ এবং কোনো সিপু কন্ট্রোল লজিক এতে জড়িত থাকে না। ইনফারেন্স ইঞ্জিনে কোনো পরিবর্তনের প্রয়োজন নেই।

ওয়েট ট্রান্সফার প্রক্রিয়াটি চার ধরনের হার্ডওয়্যার রিসোর্স ব্যবহার করতে পারে: (১) হোস্ট-ডিভাইস ডেটা মুভমেন্ট (২) প্রজেকশন ফিউশন এবং কোয়ান্টাইজেশনের জন্য জিপিইউ কম্পিউটেশন (৩) ডেটা প্লেনের জন্য RDMA নেটওয়ার্ক (৪) কন্ট্রোল প্লেনের জন্য ইথারনেট।

আমাদের ডিজাইন ওয়েট ট্রান্সফার টাস্কগুলোকে পাইপলাইন স্টেজে বিভক্ত করে, যা বিভিন্ন হার্ডওয়্যার রিসোর্স জুড়ে সহজে ওভারল্যাপিংয়ের অনুমতি দেয়।

কিছু ইমপ্লিমেন্টেশন প্রতিটি ট্রেনিং ধাপে ট্রান্সফার শিডিউল পুনরায় গণনা করে, বারবার মেটাডেটা সংগ্রহ করে এবং নির্দেশাবলী বিতরণ করে। এটি অপ্রয়োজনীয় কন্ট্রোল-প্লেন লেটেন্সি যোগ করে।

আমাদের শিডিউল ইনিশিয়ালাইজেশনের সময় একবার গণনা করা হয়। প্রতিটি ট্রেনিং ইটারেশন কেবল প্ল্যানটি পুনরায় চালায়: কন্ট্রোলার একটি “go” সংকেত জারি করে, এবং জিপিইউগুলো তাদের পূর্বনির্ধারিত রুট অনুসরণ করে। এক্সিকিউশন অনুমানযোগ্য এবং হালকা ওজনের।

সম্পূর্ণ ওয়েট আপডেট প্রক্রিয়াটিকে একটি একক মনোলিথিক ফাংশনে আটকে ফেলা লোভনীয় হতে পারে: মেটাডেটা সংগ্রহ, নাম ম্যাচিং, ইন্ট্রা-নোড গ্যাদারিং, প্রজেকশন ফিউশন, কোয়ান্টাইজেশন, সাবস্লাসিং কমিউনিকেশন ওয়ার্ল্ড, ইন্টার-নোড নেটওয়ার্ক ট্রান্সফার। এটি সঠিকভাবে প্রোগ্রাম করা কঠিন, এবং অপ্টিমাইজ করা আরও কঠিন।

আমাদের ইঞ্জিনিয়ারিংয়ে, আমরা এই পদক্ষেপগুলোকে পৃথক উপাদান হিসেবে আলাদা করি। প্রতিটি উপাদান এককভাবে টেস্ট করা যেতে পারে, বিশ্লেষণ করা যেতে পারে এবং আলাদাভাবে অপ্টিমাইজ করা যেতে পারে।

দ্রুত, নির্ভরযোগ্য ওয়েট ট্রান্সফার বড় মাপের RL ফাইন-টিউনিংয়ের জন্য একটি অত্যন্ত গুরুত্বপূর্ণ বিল্ডিং ব্লক। RDMA WRITE প্রিমিটিভ, একটি স্ট্যাটিক ট্রান্সফার শিডিউল, এবং পাইপলাইনড এক্সিকিউশন একত্রিত করে, আমরা কিমি-কে২-এ ট্রিলিয়ন-প্যারামিটার আপডেটকে মাত্র ১.৩ সেকেন্ডে কমিয়ে এনেছি। এই পদ্ধতিটি বিশ্লেষণ করা সহজ, রক্ষণাবেক্ষণ করা সহজ, এবং ট্র্যাডিশনাল ডিজাইনের বাধাগুলো এড়িয়ে চলে।

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: একমুখী ট্রান্সফার

হাই-লেভেল ওয়ার্কফ্লো

ওয়েট ট্রান্সফার এক্সিকিউশন

ডিভাইসমেশ এবং মেশ গ্রুপ

টাস্ক পাইপলাইন

জিপিইউ মেমোরি ব্যবহার নিয়ন্ত্রণ

এটি কেন দ্রুত এবং সহজ

পয়েন্ট-টু-পয়েন্ট কমিউনিকেশন

একমুখী ডেটা ট্রান্সফার

পাইপলাইনিং

স্ট্যাটিক শিডিউল

পরিচ্ছন্ন সেপারেশন

উপসংহার

উদ্ধৃতি