ما عملا اضافه کردن در زنان در ML سمپوزیوم اکتبر 19 ثبت نام:

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