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

OpenXLA یک کامپایلر مخصوص دامنه برای جبر خطی است که می‌تواند مدل‌های TensorFlow را بدون تغییر کد منبع تسریع کند.

مقدمه

هنگامی که یک برنامه 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 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]))

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

استفاده با Keras

برای مدل‌های Keras، jit_compile=True می‌توان به عنوان آرگومان برای model.compile تنظیم کرد: 

model.compile(optimizer="adam", jit_compile=True)

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

XLA:GPU را می توان با استراتژی توزیع شده TF ( MirroredStrategy یا MultiWorkerMirroredStrategy ) با حاشیه نویسی تابع step با jit_compile=True استفاده کرد: 

@tf.function(jit_compile=True)
def step_fn():
  t = tf.ones(shape=[100], dtype=tf.float32)
  ctx = tf.distribute.get_replica_context()
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)

@tf.function
def run_fn():
  return strategy.run(step_fn)

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

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

$ 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 (پیش از زمان) برای CPU با tfcompile

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

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

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

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

در ادامه مطلب

XLA Frontends

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

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

صحبت می کند

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

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