پیکره توکن‌سازی‌شده HDF5 برای آموزش مدل‌های زبانی

در پروژه‌های پیش‌آموزش مدل‌های زبانی، گلوگاه اصلی آموزش معمولاً نه توان محاسباتی پردازنده، بلکه نحوه‌ی رساندن داده‌ها از دیسک به کارت گرافیک است. اگر داده‌ها در قالب فرمت‌هایی مثل JSONL ذخیره شوند، ورکر‌های دیتالودر به دلیل کندی تجزیه متن، عدم امکان جهش مستقیم به یک اندیس خاص و سربار سنگین Page Cache، با افت شدید کارایی روبه‌رو می‌شوند و کارت گرافیک در حالت انتظار خالی (Idle) باقی می‌ماند.

فرمت HDF5 با ارائه‌ی مجموعه‌داده‌های عددی تکه‌بندی‌شده (Chunked) و قابلیت تغییر اندازه، امکان خواندن مستقیم بازه‌های توکن‌ها را بدون بارگذاری کل فایل فراهم می‌کند. در این درس یاد می‌گیرید چگونه داده‌های ورودی را حین توکن‌سازی در قطعات بهینه‌ی HDF5 بنویسید، نوشتن را روی چند فایل بخش‌بندی (Shard) کنید و با یک دیتالودر پنجره‌ی لغزان (Sliding Window)، دسته‌های ورودی و هدف را با حداکثر توان دیسک استخراج کنید.

پیکره‌ی متنی (corpus) دانلودشده باید در ساختاری ذخیره شود که موتور آموزش بتواند داده‌ها را با بالاترین سرعت خط (line speed) به‌شکل استریمینگ بخواند. فرمت JSONL روی دیسک زیر بار 16 ورکرِ دیتا لودر (dataloader) دوام نمی‌آورد؛ اما HDF5 با یک مجموعه‌داده‌ی عددی صحیح، تکه‌بندی‌شده و با قابلیت تغییر اندازه، به‌خوبی پاسخگوی این نیاز است. در این درس، فرایند توکن‌سازی استریمینگ را پیاده‌سازی می‌کنیم که خروجی را در یک مجموعه‌داده‌ی با قابلیت تغییر اندازه‌ی HDF5 می‌نویسد، داده‌ها را میان چند فایل بخش‌بندی (shard) می‌کند، در زمان آموزش خواندن داده‌ها را به‌صورت تکه‌بندی‌شده انجام می‌دهد، و در نهایت با یک دیتا لودر پنجره‌ی لغزان (sliding window)، توالی‌های با طول ثابت را همراه با بسته‌بندی مناسب تولید می‌کند.

اهداف یادگیری

• استریم کردن اسناد به یک مجموعه‌داده‌ی عددی صحیح با قابلیت تغییر اندازه در HDF5 همراه با قطعه‌بندی متن (chunking) معین و قطعی.
• بخش‌بندی نوشتن میان چند فایل HDF5 تا دامنه‌ی خرابی‌ها محدود شود و امکان اجرای موازی فراهم گردد.
• بازخوانی توکن‌ها از طریق ساختار تکه‌بندی‌شده‌ی متکی بر حافظه‌ی موقت صفحات سیستم‌عامل (page cache) در HDF5، تا دیتا لودر داده‌ها را فقط در زمان ساخت بچ به بافرهای دسته‌ای منتقل کند.
• پیاده‌سازی یک دیتا لودر پنجره‌ی لغزان که توالی‌های آموزشی با طول ثابت را طبق قواعد مشخص بسته‌بندی تولید می‌کند.

صورت مسئله

یک فرایند آموزش مدل زبانی مدرن، توکن‌ها را با نرخ صدها هزار نمونه در ثانیه در میان ده‌ها ورکر پردازش می‌کند. فایل‌های JSONL روی دیسک با اولین نقص صفحه‌ی بافر ناآماده (cold-cache page fault) از پا درمی‌آیند: پارسر JSON کند است، مرز اسناد قابل آدرس‌دهی مستقیم نیست، و رفتن به «نمونه‌ی 4,217,884» نیازمند اسکن کردن کل فایل از ابتدا است. حتی فرمت Parquet که فشرده‌سازی بسیار خوبی دارد هم انتخاب مناسبی نیست؛ چون موتور آموزش به ستون نیازی ندارد، بلکه جریانی تخت از توکن‌ها با دسترسی تصادفی $O(1)$ می‌خواهد.

فرمت HDF5 گزینه‌ی ایده‌آلی است؛ چرا که یک مجموعه‌داده‌ی تکه‌بندی‌شده، با قابلیت تغییر اندازه و منحصراً از نوع عدد صحیح ارائه می‌دهد که تکه‌های آن در زمان خواندن سازگاری کاملی با page cache دارند. موتور آموزش برشی از داده‌ها مانند tokens[3_200_000 : 3_208_192] را درخواست می‌کند و HDF5 تکه‌ی درخواستی (hyperslab) را از page cache در یک آرایه‌ی تازه تخصیص‌یافته‌ی NumPy کپی می‌کند. هزینه‌ی این کار فقط یک هندل فایل باز و فضایی به‌اندازه‌ی یک تکه از page cache برای هر ورکر است؛ هزینه‌ای که در مقایسه با رمزگشایی JSONL کاملاً ناچیز به‌شمار می‌رود.

چالش اصلی در پیاده‌سازی، نوشتن اصولی داده‌ها است. استفاده‌ی نادرست از مجموعه‌داده‌های با قابلیت تغییر اندازه بسیار آسان است: اگر اسناد را دانه‌به‌دانه بنویسید، فایل HDF5 آن‌قدر تکه‌تکه (fragmented) می‌شود که دیگر قابل استفاده نخواهد بود. اگر همه‌ی اسناد را در یک مرحله تغییر اندازه بنویسید، با متوقف شدن ناگهانی فرایند، کل آن بخش از دست می‌رود. روش اصولی این است: ابتدا بافر کردن داده‌ها و سپس افزایش اندازه، با اندازه‌ی بافری برابر با اندازه‌ی هر تکه، و همچنین نوشتن بخش‌بندی‌شده میان چند فایل تا در صورت وقوع کرش، حداکثر یک بخش از بین برود.

مفاهیم پایه

پیاده‌سازی درست HDF5 با قابلیت تغییر اندازه

مجموعه‌داده‌ی توکن‌ها با maxshape=(None,) و اندازه‌ی ثابت chunks=(chunk_size,) ایجاد می‌شود. نوشتن با بافر کردن توکن‌ها درون یک آرایه‌ی NumPy با طول chunk_size پیش می‌رود. وقتی بافر پر شد، اندازه‌ی مجموعه‌داده دقیقاً به‌میزان chunk_size افزایش می‌یابد و محتوای بافر در محدوده‌ی جدید نوشته می‌شود. در پایان هر بخش، بافر باقی‌مانده در یک محدوده‌ی نهایی و ناقص نوشته می‌شود. همه‌ی عملیات نوشتن به‌جز مورد آخر، پیوسته در حافظه و هم‌تراز با تکه‌ها هستند؛ در مورد آخر هم به خواننده اعلام می‌شود که خواندن را در نقطه‌ی token_count ثبت‌شده در ویژگی‌های (attributes) HDF5 متوقف کند.

نوشتن بخش‌بندی‌شده

یک فایل HDF5 منفرد یعنی یک نقطه‌ی شکست واحد (single point of failure). این پایپ‌لاین بخش‌ها را به‌صورت موازی می‌نویسد: هر بخش ورودی حاصل از درس 42 از فاز 19، یک بخش خروجی در فرمت HDF5 تولید می‌کند. فایلی با نام shards.json به‌عنوان نمایه، مسیر فایل، تعداد توکن‌ها، تعداد اسناد و هش sha256 محتوای توکن‌ها را برای هر بخش ثبت می‌کند. موتور آموزش با خواندن shards.json آفست‌های سراسری را محاسبه و اصالت پیکره را اعتبارسنجی می‌کند.

خواندن تکه‌بندی‌شده

در زمان آموزش، هر ورکر سهم خود را از فایل‌های HDF5 با حالت swmr=True باز می‌کند و بازه‌ی tokens[start:stop] را می‌خواند. ساختار تکه‌بندی‌شده‌ی HDF5 باعث می‌شود به‌محض آماده بودن تکه در کش، این خواندن بسیار کم‌هزینه باشد: کتابخانه‌ی h5py فایل را مستقیماً نگاشت حافظه (mmap) نمی‌کند، بلکه تکه‌های کامل را از طریق کش تکه‌های HDF5 و page cache سیستم‌عامل می‌خواند و سپس برش موردنظر را کپی می‌کند. ورکر هرگز کل فایل را در حافظه بارگذاری نمی‌کند: برش برداشته‌شده درون بافر دسته‌ای دیتا لودر کپی می‌شود، و دیتا لودر در زمان ساخت بچ آن را به یک تنسور آموزشی در حافظه‌ی پین‌شده منتقل می‌کند. مسیر اجرای سریع (hot path) فقط به ازای تغییر هر تکه یک فراخوانی سیستمی (syscall) دارد و باقی عملیات مستقیماً در حافظه‌ی رم انجام می‌شود.

دیتا لودر پنجره‌ی لغزان

دیتا لودر تنها بخشی از سیستم است که طول توالی آموزش را می‌داند. این بخش یک اندیس شروع تصادفی در جریان سراسری توکن‌ها انتخاب می‌کند، تعداد window_size + 1 توکن می‌خواند، و خروجی را به‌صورت (input, target) = (tokens[:-1], tokens[1:]) بازمی‌گرداند. مرز اسناد الزامی برای جداسازی ایجاد نمی‌کند: یک پنجره ممکن است روی دو سند هم‌پوشانی داشته باشد و یک boundary_token_id مشخص بین آن‌ها قرار گیرد تا مدل یاد بگیرد از این جداکننده استفاده کند. این یک قاعده‌ی استاندارد در بسته‌بندی است؛ همان قاعده‌ای که افراد تازه‌کار فراموش می‌کنند و در نتیجه به پیکره‌ای می‌رسند که 8 درصد آن توکن‌های مرزی و 92 درصد آن متن طبیعی است.

پیاده‌سازی از صفر

فایل code/main.py این بخش‌ها را پیاده‌سازی می‌کند:

• Tokenizer - یک توکن‌ساز معین در سطح بایت که برای این دمو کافی است. رابط آن شامل encode(text) -> list[int] و vocab_size است.
• HDF5ShardWriter - یک مجموعه‌داده‌ی عددی با قابلیت تغییر اندازه را باز می‌کند، توکن‌ها را تا رسیدن به اندازه‌ی تکه بافر می‌کند، با استرایدهای هم‌اندازه تغییر سایز داده و می‌نویسد، و هنگام بستن فایل، token_count و sha256 را در ویژگی‌های HDF5 ذخیره می‌کند.
• ShardedTokenizationPipeline - روی اسناد ورودی پیمایش می‌کند، آن‌ها را به یک نویسنده می‌سپارد، و نمایه‌ی shards.json را تولید می‌کند.
• MmapTokenStore - فایل‌های بخش‌بندی‌شده را برای خواندن تکه‌بندی‌شده باز می‌کند (این نام صرفاً جهت سهولت انتخاب شده؛ h5py تکه‌ها را می‌خواند و فایل را mmap نمی‌کند)، آفست‌های سراسری را محاسبه می‌کند، و یک API واحد به‌صورت get_slice(start, stop) در اختیار می‌گذارد.
• SlidingWindowDataloader - پنجره‌های تصادفی از جریان سراسری برمی‌دارد و آرایه‌های NumPy شامل (input_ids, target_ids) تولید می‌کند.

یک دموی اجرایی در انتهای فایل، پیکره‌ی متنی بسیار کوچکی در حافظه می‌سازد، آن را درون دو بخش توکن‌سازی می‌کند، فایل‌ها را برای خواندن تکه‌بندی‌شده باز می‌کند، دیتا لودر را برای 10 بچ اجرا می‌نماید، و شکل ابعاد و چکسام (checksum) هر بچ را چاپ می‌کند.

برای اجرا:

python3 code/main.py

اسکریپت با کد خروجی صفر به پایان می‌رسد و چکسام‌های بچ‌ها را در خروجی چاپ می‌کند.

الگوهای محیط عملیاتی

چهار الگو به شما کمک می‌کنند این درس را در یک فرایند آموزش واقعی در محیط عملیاتی پیاده کنید.

اندازه‌ی تکه برابر با اندازه‌ی خواندن معمول باشد. موتور آموزش برای هر نمونه window_size + 1 توکن می‌خواند. اگر اندازه‌ی تکه‌ی HDF5 مضربی از window_size باشد، خواندن‌ها هم‌راستا با page cache انجام می‌شوند. در صورت عدم تطابق این اندازه‌ها، نرخ پردازش نصف می‌شود؛ چرا که هر نمونه دو تکه‌ی متفاوت را درگیر خواهد کرد.

تعداد توکن‌ها در ویژگی‌ها ثبت شود، نه در ابعاد مجموعه‌داده. آخرین برش مجموعه‌داده ممکن است ناقص پر شده باشد، زیرا اندازه‌ی تکه لزوماً بر مرز اسناد بخش‌پذیر نیست. مقدار واقعی token_count را در قالب ویژگی‌های HDF5 روی مجموعه‌داده ذخیره کنید و از خواننده بخواهید داده‌ها را تا همین مقدار بخواند. در غیر این صورت، خواننده از مرز داده عبور کرده و توکن‌های پرشده با صفر را می‌خواند؛ در نتیجه مدل یاد می‌گیرد مقدار صفر را پیش‌بینی کند.

هش sha256 مجزا برای هر بخش همراه با اعتبارسنجی موازی. هر بخش sha256 اختصاصی خود را روی بایت‌های توکن‌ها دارد. موتور آموزش می‌تواند همه‌ی بخش‌ها را پیش از شروع آموزش به‌صورت موازی بررسی و اعتبارسنجی کند. اگر مقدار sha256 نادرست باشد، فرایند در همان ابتدا متوقف می‌شود، نه در دور (epoch) سوم و بعد از 16 ساعت کار.

تنظیم swmr=True در هر دو طرف همراه با libver="latest" در نویسنده. حالت یک‌نویسنده-چندخواننده (Single-Writer-Multiple-Reader یا SWMR) ایجاب می‌کند که نویسنده فایل را با libver="latest" باز کند، همه‌ی مجموعه‌داده‌ها را از قبل بسازد و سپس file.swmr_mode = True را فعال کند. پس از آن، نویسنده باید بعد از هر بار تغییر اندازه دستور dataset.flush() را فراخوانی کند تا ورکر‌های خواننده (که با swmr=True باز شده‌اند) داده‌ها را به‌صورت پایدار ببینند. نادیده گرفتن libver="latest" یا فعال کردن SWMR بعد از اعمال تغییرات ساختاری، از دلایل رایج خطای «file is locked» است.

استفاده از ابزارهای واقعی

الگوهای محیط عملیاتی (پروداکشن):

• یک فایل HDF5 به ازای هر بخش مبدأ. دانلودکننده (درس 42) به ازای هر URL یک بخش تولید می‌کند؛ توکن‌سازی (این درس) به ازای هر بخش ورودی یک فایل HDF5 می‌سازد. این نگاشت 1:1 ادامه‌ی کار پس از توقف و بازیابی پس از خرابی‌های جزئی را بسیار ساده می‌کند.
• شناسه‌ی توکن مرزی. توکن مرزی جزئی از دایره‌ی واژگان توکن‌ساز است و تنها توکنی است که دیتا لودر به داده‌ها اضافه می‌کند. اگر قرار است مدل این توکن را نادیده بگیرد، در تابع زیان (loss) ماسک می‌شود؛ در غیر این صورت مدل یاد می‌گیرد از آن به‌عنوان جداکننده‌ی توالی استفاده کند.
• فایل shards.json به‌عنوان منبع حقیقت. افزودن یک بخش جدید یعنی نوشتن فایل HDF5، محاسبه‌ی sha256 آن، و اضافه کردن یک رکورد جدید به این فایل. موتور آموزش فقط یک‌بار در ابتدای کار این فایل را می‌خواند و هیچ‌وقت نیازی به بررسی فهرست فایل‌های دایرکتوری نخواهد داشت.

خروجی این درس

فایل outputs/skill-hdf5-tokenized-corpus.md در یک پروژه‌ی واقعی توضیح می‌دهد که کدام توکن‌ساز ورودی پایپ‌لاین را تأمین می‌کند، چه اندازه‌ی تکه‌ای با پنجره‌ی آموزش تطابق دارد، فایل shards.json در کجای سیستم کنترل نسخه (version control) نگهداری می‌شود، و ورکر‌های دیتا لودر چگونه روی فایل‌ها تقسیم کار می‌کنند. در این درس، موتور اصلی این سازوکار پیاده‌سازی و ارائه شده است.

تمرین‌ها

1. پرچم --compression gzip را به نویسنده‌ی HDF5 اضافه کنید و افت نرخ پردازش را روی پیکره‌ی نمایشی بسنجید. سپس در مورد گزینه‌ی پیش‌فرضی که ترجیح می‌دهید، استدلال کنید.
2. یک سید (seed) معین به دیتا لودر پنجره‌ی لغزان اضافه کنید و بررسی کنید که آیا دو بار اجرای متوالی با همان سید، دسته‌های آموزشی کاملاً یکسانی تولید می‌کنند یا خیر.
3. یک حالت --validate بسازید که همه‌ی بخش‌ها را بخواند، sha256 توکن‌های آن‌ها را دوباره محاسبه کند و آن را با مقادیر shards.json تطبیق دهد. پایپ‌لاین CI باید قبل از شروع آموزش این بررسی را اجرا کند.
4. نرخ پردازش دیتا لودر را با اندازه‌ی تکه‌های برابر با، نصف، و دو برابرِ اندازه‌ی پنجره مقایسه کنید. اثر page cache را در قالب یک گزارش تحلیل کنید.
5. پرچم --max-document-tokens را اضافه کنید تا اسناد خیلی طولانی را در زمان نوشتن کوتاه کند. مزایا و معایب این کار را در مقایسه با اعمال آن در زمان خواندن تحلیل کنید.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• مستندات تکه‌بندی در HDF5 - آشنایی با ساختار مجموعه‌داده‌های تکه‌بندی‌شده و قابل تغییر اندازه که در این درس استفاده شد
• راهنمای کاربری h5py - اتصالات پایتون برای کار با HDF5
• ذخیره‌سازی تکه‌بندی‌شده در h5py - نحوه‌ی خواندن مجموعه‌داده‌های تکه‌بندی‌شده توسط h5py و تنظیم اندازه‌ی کش تکه‌ها
• فاز 19 · درس 42 - ماژول دانلودکننده‌ای که خروجی آن در این درس توکن‌سازی می‌شود
• فاز 19 · درس 44 - برنامه‌ی زمان‌بندی کسینوسی (cosine schedule) که خروجی این دیتا لودر را مصرف می‌کند
• فاز 19 · درس 45 - حلقه‌ی دقت ترکیبی خودکار (AMP) که گام‌های آموزش را در بر می‌گیرد