تنظیم پرچم‌های مگااسکیل XLA

MegascaleXLA (MXLA) ارتباطات بین برش‌های TPU را از طریق شبکه‌های مرکز داده (DCN) هماهنگ می‌کند. این راهنما دستورالعمل‌ها و بهترین شیوه‌ها را برای پیکربندی و تنظیم پرچم‌های MegascaleXLA برای کاربرانی که بارهای کاری TPU چند برشی را اجرا می‌کنند، ارائه می‌دهد. تلاش شده است تا برای اکثر موارد، پیش‌فرض‌های معقولی تنظیم شود.

اندازه‌های بافر جمعی

برای کالکتیوهای کوچک، عموماً بسته‌بندی ارسال/دریافت‌های TPU کوچک و ارسال/دریافت‌های شبکه در کالکتیوهای بزرگتر برای کاهش سربار تأخیر، کارآمدتر است. برای کالکتیوهای بزرگتر، استفاده از تعداد بیشتری از بافرهای میانی کوچکتر ممکن است مفید باشد. تعدادی ابزار وجود دارد که امکان تغییر این رفتار را فراهم می‌کنند. خواندن‌های TPU DMA را می‌توان برای برخی از کالکتیوها با عبور دادن، ادغام یا تقسیم کرد.

--megascale_target_dma_size=<value>

که تلاش می‌کند تا خواندن‌های DMA را تا حد امکان به اندازه هدف نزدیک کند. به طور پیش‌فرض، از مقدار پیش‌فرض ۸ مگابایت استفاده می‌شود. فقط توصیه می‌شود این مقدار را تنظیم کنید اگر ردیابی پروفایل، تأخیر قابل توجهی را بین زمان شروع DMA و شروع اولین انتقال شبکه برای یک مجموعه خاص نشان دهد.

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

--megascale_sparse_reduction_threshold=<value>

به طور پیش‌فرض، این مقدار ۲۵۶ کیلوبایت است. اگر یک پروفایلر سربار تأخیر قابل توجهی را برای عملیات کاهش نشان دهد، ممکن است تنظیم این مقدار به سمت بالا مفید باشد تا از کاهش‌دهنده‌های کمتری استفاده شود و در نتیجه تعداد کمتری از انتقال‌های بزرگتر در شبکه ایجاد شود.

در برخی شرایط، ممکن است لازم باشد از انتقال‌های بزرگ شبکه اجتناب شود و در عوض از تعداد بیشتری انتقال کوچک‌تر استفاده شود. این مورد را می‌توان با استفاده از پرچم تنظیم کرد:

--megascale_max_reduction_shard_size=<value>

که بافر را به تعداد بیشتری از قطعات کوچکتر از آنچه که در غیر این صورت توسط مجموعه مورد نیاز است تقسیم می‌کند. به طور پیش‌فرض، این مقدار ۸ مگابایت است. باید به این صورت باشد که megascale_max_reduction_shard_size >> megascale_sparse_reduction_threshold زیرا این دو اثرات متضادی دارند.

برای انتقال‌های یک به یک (مانند collective-permute )، انتقال‌ها روی یک میزبان را می‌توان با استفاده از موارد زیر به تکه‌هایی تقسیم کرد:

--megascale_chunk_size=<value>

اگر صفر نباشد، این هر انتقال یک به یک روی یک میزبان را به چندین تکه بایت megascale_chunk_size (پیش‌فرض: ۸ مگابایت) تقسیم می‌کند تا انتقال‌های DMA و شبکه را به صورت پایپ‌لاین انجام دهد. همه کالکتیوهای دیگر از megascale_target_dma_size ، megascale_sparse_reduction_threshold و megascale_max_reduction_shard_size برای کنترل خرد کردن و قطعه‌بندی استفاده می‌کنند.


مدیریت حافظه و انتقال‌های بدون کپی

انتقال‌های شبکه در مقیاس بزرگ به گونه‌ای طراحی شده‌اند که کپی صفر باشند: نواحی حافظه میزبان برای DMA به و از TPU پین و از پیش نگاشت می‌شوند. اگر حافظه از پیش نگاشت شده در زمان اجرا تمام شود، سیستم به نگاشت بر اساس تقاضا یا کپی‌های حافظه پویا بازمی‌گردد که به شدت بر توان عملیاتی تأثیر می‌گذارد.

ناحیه حافظه از پیش نگاشت شده

--megascale_grpc_premap_memory_bytes=17179869184  # Default: 16 GB (16LL << 30)

اندازه ناحیه حافظه پین‌شده میزبان که در زمان مقداردهی اولیه برای انتقال‌های DMA اختصاص داده شده است را کنترل می‌کند.

علائم تشخیصی

  1. MapDmaBuffer در حالت پایدار : در XProf Trace Viewer، MapDmaBuffer را جستجو کنید. این فراخوانی‌ها فقط باید در هنگام راه‌اندازی رخ دهند. اگر در طول تکرارهای آموزش ادامه یابند، نگاشت پویا در حال رخ دادن است.
  2. ردیابی‌های Megascale: Memory Copy : اگر مسیر انتقال ارتباطات در XProf رویدادهای Megascale: Memory Copy را نمایش دهد، بافرهای حافظه به جای انتقال از طریق DMA با کپی صفر، کپی می‌شوند.

وضوح تصویر

مقدار --megascale_grpc_premap_memory_bytes افزایش دهید (مثلاً به ۲۴ گیگابایت یا ۳۲ گیگابایت، بسته به حافظه میزبان موجود در نمونه ماشین مجازی TPU شما) و کار را مجدداً راه اندازی کنید.