XLA: بهینه سازی کامپایلر برای یادگیری ماشین

XLA (Accelerated Linear Algebra) یک کامپایلر مخصوص دامنه برای جبر خطی است که می تواند مدل های TensorFlow را بدون تغییر کد منبع شتاب دهد.

به عنوان مثال در برت: نتایج بهبود در سرعت و استفاده از حافظه هستند MLPerf تسلیم با استفاده از 8 ولتا V100 GPU ها با استفاده از XLA است ~ 7X بهبود عملکرد بهبود اندازه ~ 5x به دسته ای به دست آورد و:

معرفی

هنگامی که یک برنامه TensorFlow اجرا می شود، تمام عملیات ها به صورت جداگانه توسط مجری TensorFlow اجرا می شوند. هر عملیات TensorFlow دارای یک هسته GPU از پیش کامپایل شده است که مجری آن را به آن ارسال می کند.

XLA یک حالت جایگزین برای اجرای مدل‌ها ارائه می‌کند: گراف TensorFlow را به دنباله‌ای از هسته‌های محاسباتی که به‌طور خاص برای مدل داده‌شده تولید می‌شوند، کامپایل می‌کند. از آنجایی که این هسته ها منحصر به مدل هستند، می توانند از اطلاعات خاص مدل برای بهینه سازی بهره برداری کنند. به عنوان مثال، اجازه دهید به بهینه سازی XLA در زمینه یک محاسبه ساده TensorFlow نگاه کنیم:

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

بدون XLA اجرا شود، گراف سه هسته را راه اندازی می کند: یکی برای ضرب، یکی برای جمع و دیگری برای کاهش. با این حال، XLA می‌تواند نمودار را طوری بهینه کند که نتیجه را در یک کرنل محاسبه کند. این کار را با "ادغام" جمع، ضرب و کاهش در یک هسته GPU انجام می دهد. علاوه بر این، این عملیات ذوب می کند ارسال نمی کردن نتایج متوسط تولید شده توسط y*z و x+y*z به حافظه؛ در عوض نتایج این محاسبات میانی را مستقیماً به کاربران خود ارسال می کند در حالی که آنها را کاملاً در رجیسترهای GPU نگه می دارد. Fusion مهمترین بهینه سازی XLA است. پهنای باند حافظه معمولاً کمیاب ترین منبع در شتاب دهنده های سخت افزاری است، بنابراین حذف عملیات حافظه یکی از بهترین راه ها برای بهبود عملکرد است.

XLA را برای مدل های TensorFlow فعال کنید

تلفیقی صریح با tf.function(jit_compile=True)

API کامپایل صریح یک کنترل دقیق برای انتخاب توابع کامپایل ارائه می دهد. به عنوان مثال، تابع TensorFlow زیر که آموزش MNIST را انجام می دهد با XLA کامپایل شده است:

@tf.function(jit_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))

jit_compile API معنای باید کامپایل: یا کل تابع با XLA وارد شده، یا یک errors.InvalidArgumentError استثنا پرتاب می شود. این است که اگر این امکان وجود ندارد برای پی بردن به ابعاد تمام تانسورها بدون در حال اجرا کل محاسبه: XLA نمی تواند در حال حاضر توابع که در آن ابعاد inferrable نمی کامپایل. به عنوان مثال، تابع زیر کامپایل نخواهد شد:

@tf.function
def not_compilable(x):
  return tf.unique(x)

شکل ها می توانند در طول اجراها متفاوت باشند:

@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))

مراجعه کنید COLAB آموزش برای یک مثال استفاده دقیق تر، و یک ویدیو آموزش در jit_compile=True استفاده.

خوشه بندی خودکار

یک راه ساده برای شروع به استفاده از XLA در مدل های TensorFlow بدون هیچ گونه تغییر است برای فعال کردن خودکار خوشه، که به طور خودکار پیدا خوشه (زیرگرافهای متصل) در توابع TensorFlow است که می تواند وارد شده و اجرا با استفاده از XLA. خودکار خوشه در GPU می تواند با تنظیم فعال TF_XLA_FLAGS متغیر محیطی:

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

خودکار خوشه در حال حاضر برای حجم کار پردازنده گرافیکی بهینه سازی شده، اما آن را نیز می توانید بر روی CPU فعال باشد توسط علاوه بر این با استفاده از پرچم --tf_xla_cpu_global_jit :

$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

برای یک مثال استفاده دقیق را ببینید خودکار خوشه آموزش COLAB .

AOT (جلوتر از زمان) تلفیقی برای CPU با tfcompile

شما همچنین می توانید مستقل استفاده کنید tfcompile ابزار، که تبدیل نمودار TensorFlow به کد اجرایی (فقط برای پردازنده x86-64 عمل).

برنامه های کامپایل شده را بررسی کنید

XLA امکانات درون نگری را فراهم می کند که به شما امکان می دهد برنامه های تولید شده را بررسی کنید. به روگرفت برنامه تولید، با استفاده از متغیر محیطی XLA_FLAGS :

$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

پس از آزادسازی انجام می شود، شما می توانید فایل زیر را پیدا کنید در /tmp/generated :

  • module_XXXX.*_optimizations.txt تولید برنامه های وابسته به جنس ، یکی در هر هر خوشه وارد شده است. پیوست کردن آن‌ها هنگام ارسال گزارش‌های باگ XLA بسیار مفید است!

  • module_XXXX.ir-*.ll فایل های تولید شده در LLVM نمایش میانی، با NVPTX intrinsics.

  • module_XXXX.ptx تولید PTX فایل های.

همچنین می توانید نمودار را با تجسم تعبیه خوشه های XLA در داخل نمودار TensorFlow با استفاده از:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

گزارش های باگ قابل تکرار

بازتولید گزارش اشکال بسیار آسان‌تر است اگر شامل موارد خالی برای برنامه‌های XLA تولید شده و تعبیه خوشه‌بندی خودکار مورد استفاده باشد. برای تولید آنها برای یک برنامه TensorFlow که با خوشه بندی خودکار اجرا می شود، راه اندازی کنید:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"

هنگام تشکیل پرونده اشکالات، ضمیمه محتویات /tmp/generated دایرکتوری (اشاره شده در بالا).

در صورت امکان، سعی کنید برای منزوی کردن مشکلات را به یک برنامه XLA تنها با استفاده از replay_computation و به صورت تکراری اجرای آن بر روی برنامه تولید می شود.

خواندن بیشتر

XLA Frontends

به غیر از TensorFlow، برنامه های XLA را می توان با موارد زیر تولید کرد:

  • JAX تحولات ساخت مجموع از برنامه های پایتون + نامپای:
  • جولیا : زبان جولیا برای محاسبات علمی
  • PyTorch چارچوب PyTorch:
  • NX : عددی کتابخانه محاسبات برای زبان برنامه نویسی اکسیر

صحبت می کند

با استفاده از XLA از TF با استفاده از jit_compile=True

بررسی اجمالی XLA