این صفحه به‌وسیله ‏Cloud Translation API‏ ترجمه شده است.
Switch to English

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

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

نتایج به دست آمده بهبودهایی در سرعت و میزان استفاده از حافظه است: به عنوان مثال در ارسال BERT MLPerf با استفاده از 8 GPU ولتا V100 با استفاده از XLA به بهبود عملکرد 7x and و بهبود اندازه دسته 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))

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

@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]))

برای مثال دقیق تر ، به کولاب آموزش مراجعه کنید.

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

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

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

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

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

برای یک مثال کاربرد دقیق ، به کلاک آموزش خوشه بندی خودکار مراجعه کنید .

تدوین AOT (پیش از زمان) برای پردازنده با 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 تولید شده ، از هر خوشه کامپایل شده ، یکی. پیوستن به کسانی که گزارشهای اشکال XLA را ارائه می دهند بسیار مفید است!

  • module_XXXX.ir-*.ll پرونده های تولید شده در نمایندگی متوسط LLVM ، با ذاتیات NVPTX .

  • 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 ضمیمه کنید (که در بالا به آنها اشاره شده است).

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

بیشتر خواندن

XLA Frontends

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

  • JAX : تحولات ترکیبی برنامه های Python + NumPy
  • جولیا : زبان جولیا برای محاسبات علمی
  • PyTorch : چارچوب PyTorch