FlashREINFORCE کے ساتھ ایجنٹک لینگویج ماڈلز کی تربیت

FlashREINFORCE کیا ہے؟

FlashREINFORCE ایجنٹ پر مبنی زبان کے ماڈلز کی غیر متزامن تربیت کے لیے ایک ناقد سے پاک reinforcement-learning فریم ورک ہے۔ یہ طویل مدتی کاموں کے لیے بنایا گیا ہے، جن میں متغیر طوالت کی reasoning، ٹول کالز، یا ماحول کے ساتھ تعامل کی وجہ سے trajectories غیر باقاعدہ اوقات میں مکمل ہوتی ہیں۔

ہر prompt کے لیے کئی ہم سطح rollouts تیار کرنے کے بجائے، FlashREINFORCE ہر prompt کے لیے ایک rollout استعمال کرتا ہے۔ اس لیے B trajectories پر مشتمل batch میں B مختلف prompts شامل ہو سکتے ہیں، اور مکمل ہونے والی trajectories کو ایک ہی prompt کے لیے متعدد جوابات کا انتظار کیے بغیر learner تک پہنچایا جا سکتا ہے۔

repository میں ایک خودمختار PyTorch reference loss، CPU optimizer-update کی مثال، tests، R1 اور Qwen2.5-Math کے لیے Molt launchers، اور reasoning، Python tools، mixture-of-experts models، اور ALFWorld کے لیے تجرباتی settings شامل ہیں۔ مکمل asynchronous training launchers Molt استعمال کرتے ہیں۔

Reinforcement Learning کو REINFORCE ہی کرنا چاہیے۔

ایک rollout کیوں استعمال کریں؟

Group-relative طریقے ایک ہی prompt سے متعدد rollouts پر منحصر ہوتے ہیں۔ غیر متزامن agent workloads میں یہ تقاضا ہم وقت سازی کی رکاوٹیں پیدا کر سکتا ہے اور مقررہ rollout budget میں شامل مختلف prompts کی تعداد کم کر سکتا ہے۔

Single-rollout learning ہم سطح rollout کی ضرورت ختم کرتی ہے، لیکن اس سے استحکام کے تین مسائل پیدا ہوتے ہیں:

• ایسا کوئی ہم سطح group نہیں ہوتا جس سے prompt-level baseline بنایا جا سکے۔
• غیر متزامن طور پر جمع کی گئی trajectories شاید کسی پرانی behavior policy نے تیار کی ہوں۔
• طویل ناکام trajectories کسی update پر غالب آ سکتی ہیں، اگر ہر token کو batch-level وزن یکساں دیا جائے۔

FlashREINFORCE ان مسائل کو One-Batch REINFORCE، Sequence Trust Region، اور Sample-Mean Optimization کے ذریعے حل کرتا ہے۔

بنیادی اجزا

1. One-Batch REINFORCE

ہر نئے batch کے لیے learner آزاد prompts کے درمیان scalar rewards کو مرکز میں لاتا ہے۔ اگر R_i trajectory i کا reward ہے، تو اس کا advantage trajectory reward اور batch mean کا فرق ہوتا ہے:

mean_reward = sum(R_j for j in batch) / B
A_i = R_i - mean_reward

اوسط سے زیادہ والی trajectories کو مثبت advantages ملتے ہیں، جبکہ اوسط سے کم والی trajectories کو منفی advantages ملتے ہیں۔ اس طرح learned critic کی تربیت کے بغیر علامتی feedback فراہم ہوتا ہے۔ learner batch پر ایک optimizer update کرتا ہے اور پھر اسے ضائع کر دیتا ہے۔

2. Sequence Trust Region

ہر trajectory کو وہ behavior-policy probability محفوظ رکھنی چاہیے جس نے حقیقت میں ہر sampled token تیار کیا تھا۔ learner موجودہ policy probability دوبارہ شمار کرتا ہے اور token importance ratio بناتا ہے:

rho_i_t = exp(current_log_probability_i_t
              - behavior_log_probability_i_t)

یہ ratio محفوظ شدہ history پر sampled action distribution کی اصلاح کرتا ہے۔ یہ اس حقیقت کی اصلاح نہیں کرتا کہ history خود کسی پرانی behavior policy سے آئی تھی، اس لیے FlashREINFORCE مکمل trajectory پر جمع شدہ drift کا بھی جائزہ لیتا ہے۔

ہر sampled action کے لیے یہ behavior probability p اور موجودہ probability q سے Bernoulli KL proxy شمار کرتا ہے:

d_i_t = p * log(p / q)
        + (1 - p) * log((1 - p) / (1 - q))

token values کا اوسط trajectory پر لیا جاتا ہے۔ کسی trajectory کو صرف اس وقت قبول کیا جاتا ہے جب اس کا mean proxy مقررہ trust threshold سے زیادہ نہ ہو:

mean_drift_i = sum(d_i_t for t in trajectory_i) / T_i
mask_i = 1 if mean_drift_i <= delta else 0

مسترد کی گئی trajectory کسی token loss میں حصہ نہیں ڈالتی۔ اس طرح trust کا فیصلہ الگ الگ tokens کے بجائے مکمل محفوظ شدہ history پر لاگو ہوتا ہے۔

3. Sample-Mean Optimization

Token-level batch average طبعی طور پر طویل trajectories کو زیادہ اثر و رسوخ دیتا ہے۔ یہ نقصان دہ ہو سکتا ہے جب ایک طویل ناکام rollout میں مفید درمیانی مراحل شامل ہوں لیکن اختتام منفی reward پر ہو۔

FlashREINFORCE پہلے ہر trajectory کے اندر token contributions کا اوسط لیتا ہے، پھر trajectories کے درمیان اوسط لیتا ہے:

trajectory_term_i = (
    mask_i * advantage_i / trajectory_length_i
    * sum(importance_ratio_i_t for each token t)
)
objective = sum(trajectory_term_i for i in batch) / B

نتیجتاً ہر trajectory کو response length سے قطع نظر 1/B وزن ملتا ہے۔ عملی loss detached importance ratios استعمال کرتا ہے۔ بنیادی طریقہ learned critic، ratio clipping، یا reference-model forward pass استعمال نہیں کرتا۔

Installation اور مقامی setup

FlashREINFORCE repository کو clone یا download کریں، اس کی root directory میں جائیں، اور test dependencies کے ساتھ اسے editable mode میں install کریں:

pip install -e '.[test]'

اس کے بعد شامل CPU example چلائیں:

python examples/toy_update.py

یہ example جانچتا ہے کہ reference objective CPU پر optimizer update پیدا کر سکتا ہے۔ یہ functional verification example ہے، paper کے benchmark نتائج کی reproduction نہیں۔

آخر میں test suite چلائیں:

python -m pytest -q

example اور tests کا کامیاب run اس بات کی تصدیق کرتا ہے کہ مقامی reference implementation اور اس کی dependencies دستیاب ہیں۔

Reference implementation کا جائزہ

مرکزی reference objective flashreinforce/loss.py میں موجود ہے۔ یہ paper اور README میں بیان کردہ operations کا احاطہ کرتا ہے:

• Batch کے درمیان reward centering۔
• Token-level importance sampling۔
• Sequence-level trust admission۔
• Sample-mean optimization۔
• Failure tokens کی entropy-based filtering، جو اختیاری ہے۔

یہ دیکھنے کے لیے کہ reference loss optimizer update میں کیسے شامل ہوتا ہے، examples/toy_update.py سے شروع کریں۔ پھر flashreinforce/loss.py کا جائزہ لیں تاکہ example tensors اور masks کو objective کے انفرادی مراحل سے مربوط کیا جا سکے۔

CPU کی مثال بنیادی استعمال کا سب سے محفوظ طریقہ ہے، کیونکہ اس کے لیے مکمل تقسیم شدہ رول آؤٹ انفراسٹرکچر درکار نہیں ہوتا۔

ایک مرحلے کی غیر ہم وقت اپ ڈیٹ کو سمجھیں

پروڈکشن FlashREINFORCE اپ ڈیٹ چھ مراحل پر مشتمل ہوتی ہے:

1. جمع کرنا: رول آؤٹ ورکرز ہر پرامپٹ کے لیے آزادانہ طور پر ایک مکمل trajectory جمع کراتے ہیں۔ ہر trajectory میں اس کا reward اور tokens کو sample کرنے کے لیے استعمال ہونے والے حقیقی behavior-policy probabilities شامل ہوتے ہیں۔
2. بیچ بنانا: learner اگلی B مکمل trajectories لیتا ہے اور signed advantages حاصل کرنے کے لیے ان کے rewards کو مرکزیت دیتا ہے۔
3. درست کرنا: learner موجودہ policy کے log probabilities دوبارہ شمار کرتا ہے اور token importance ratios بناتا ہے۔
4. اعتماد: یہ ہر trajectory کے لیے sampled-action KL proxy کا اوسط شمار کرتا ہے اور مکمل trajectory کو صرف اس صورت قبول کرتا ہے جب قدر مقررہ threshold کے اندر ہو۔
5. بہتر بنانا: یہ ہر منظور شدہ trajectory کے اندر اوسط لیتا ہے، پھر پوری batch میں اوسط لیتا ہے، اور optimizer کا ایک قدم اٹھاتا ہے۔
6. خارج کرنا: یہ اضافی learner updates کے لیے جمع کی گئی batch کو دوبارہ استعمال نہیں کرتا۔

آخری خارج کرنے کا مرحلہ اس طریقے کا اہم حصہ ہے۔ FlashREINFORCE تازہ جمع کی گئی ہر batch کے لیے ایک full-batch update استعمال کرتا ہے، بجائے اس کے کہ اسی behavior snapshot سے کئی مسلسل minibatch updates چلائے جائیں۔

Molt کی training configuration کا پیش نظارہ

repository میں scripts/train_molt.py شامل ہے، جو Molt کے asynchronous implementation کے ذریعے training شروع کرنے کے لیے ہے۔ مکمل training run کے لیے وسائل مختص کرنے سے پہلے، training flags کا پیش نظارہ کرنے کے لیے اس کا dry-run mode استعمال کریں:

python scripts/train_molt.py --dry-run

GPU setup، data preparation، training commands، paper سے code تک mapping، اور معلوم reproduction limitations کے لیے docs/training.md سے رجوع کریں۔ examples/README.md گائیڈ میں task presets، ablations، اور capability checks کے ساتھ native trainer launcher شامل ہے۔

R1 اور Qwen2.5-Math launchers repository میں pinned ہیں۔ Tool-use اور ALFWorld settings کے لیے compatible agents اور environments بھی درکار ہیں، اس لیے ان experiments کو شروع کرنے سے پہلے ان کی runtime requirements کا جائزہ لیں۔

وہ data جو ہر trajectory کو محفوظ رکھنا چاہیے

یہ update generation کے دوران حاصل کی گئی معلومات پر منحصر ہے۔ کم از کم، conceptual training flow کو مکمل token sequence، اس کا scalar reward، اور ہر sampled action سے وابستہ behavior probability درکار ہوتی ہے۔ اس کے بعد learner محفوظ شدہ actions اور histories کا جائزہ لے کر متعلقہ current-policy probabilities حاصل کرتا ہے۔

ریکارڈ شدہ sampling probabilities کو بعد میں دوبارہ تشکیل دی گئی probabilities سے تبدیل نہ کریں۔ بعد میں کیا گیا model evaluation اس probability کو دوبارہ پیدا کرنے کی ضمانت نہیں دیتا جسے inference engine نے sampling کے وقت حقیقتاً استعمال کیا تھا۔

استحکام کے لیے جدید نکات

Sampling support برقرار رکھیں

درست importance correction یہ فرض کرتی ہے کہ target policy، behavior policy کے support سے باہر probability mass مختص نہیں کرتی۔ جارحانہ top-k یا top-p truncation اس مفروضے کی خلاف ورزی کر سکتی ہے۔ اس لیے sampling settings منتخب کرتے وقت support requirement کو مدنظر رکھنا چاہیے۔

Token importance sampling کو مکمل trajectory correction نہ سمجھیں

Token importance ratios، محفوظ شدہ histories کی بنیاد پر sampled actions کی correction کرتے ہیں۔ وہ ان histories کی distribution کو درست نہیں کرتے۔ ایسی trajectories کو filter کرنے کے لیے Sequence Trust Region استعمال کریں جن میں accumulated behavior-to-learner drift بہت زیادہ ہو۔

Updates کو تازہ رکھیں

ہر جمع کی گئی batch کے لیے optimizer کا ایک update کریں اور پھر اسے خارج کر دیں۔ پرانی batch کو کئی learner updates کے لیے دوبارہ استعمال کرنا documented one-pass method سے ہٹ جاتا ہے اور policy mismatch بڑھا سکتا ہے۔

Trajectory-level weighting برقرار رکھیں

Batch میں trajectories کا اوسط لینے سے پہلے ہر trajectory کے اندر token contributions کا اوسط لیں۔ تمام valid tokens کو ایک ہی batch-wide mean میں flatten کرنے سے response-length multiplier دوبارہ شامل ہو جائے گا اور طویل trajectories کو غیر متناسب وزن ملے گا۔

Trust mask کی درست تشریح کریں

Sequence Trust Region پوری trajectory کے لیے داخلے کا ایک فیصلہ کرتی ہے۔ اگر کوئی trajectory threshold سے تجاوز کرے، تو اس update سے اس کے تمام token losses نکال دیے جاتے ہیں۔

Verification اور reproduction کو الگ رکھیں

CPU مثال صرف optimizer update کی verification کرتی ہے۔ تحقیقی نتائج کو دوبارہ پیدا کرنے کے لیے training اور examples guides میں درج GPU، data، model، agent، environment، اور asynchronous runtime configurations درکار ہوتی ہیں۔

رپورٹ کیے گئے اطلاقی شعبے

رپورٹ کیے گئے experiments میں mathematical reasoning، multi-turn Python tool use، mixture-of-experts training، اور interactive ALFWorld decision making شامل ہیں۔ README میں DeepSeek-R1-Distill-Qwen-1.5B کے لیے 6,000 updates تک stable long-chain-of-thought training، Qwen2.5-Math-1.5B کے لیے mathematical reasoning results، Qwen2.5-7B-Instruct اور Qwen3-30B-A3B کے ساتھ Python-tool experiments، اور Qwen2.5-7B-Instruct کے ساتھ ALFWorld experiments رپورٹ کیے گئے ہیں۔

یہ نتائج framework کے مطلوبہ دائرۂ کار کو ظاہر کرتے ہیں، لیکن چھوٹی CPU update مثال سے ان کی توقع نہیں کی جانی چاہیے۔

مفید وسائل

• FlashREINFORCE ریپوزٹری
• پروجیکٹ کی ویب سائٹ
• NVIDIA NeMo Molt کوڈ
• تربیتی رہنما
• تجرباتی ترتیبات اور مثالیں

نتیجہ

FlashREINFORCE سنگل رول آؤٹ پرامپٹ کوریج کو ایک مختصر، نقاد سے آزاد اپ ڈیٹ کے ساتھ یکجا کرتا ہے۔ بیچ پر مبنی انعامات علامت دار ایڈوانٹیجز فراہم کرتے ہیں، ٹوکن امپورٹنس سیمپلنگ پرانے اعمال سے نمٹتی ہے، Sequence Trust Region جمع شدہ ٹریجیکٹری ڈرفٹ کو کنٹرول کرتا ہے، اور سیمپل میان آپٹیمائزیشن ردِعمل کی لمبائی کو ٹریجیکٹری کے وزن کا تعین کرنے سے روکتی ہے۔

قابلِ ترمیم انسٹالیشن، CPU اپ ڈیٹ کی مثال اور ٹیسٹس سے آغاز کریں۔ ان کے کام کرنے کے بعد، --dry-run کے ساتھ Molt لانچر کا پیش منظر دیکھیں اور ماڈل کے لحاظ سے GPU اور ڈیٹا کی تیاری کے لیے تربیتی رہنما کی پیروی کریں۔ یہ پروجیکٹ Apache License 2.0 کے تحت دستیاب ہے۔