راهنمای کاربر جداسازی HLO

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

نصب

شما می‌توانید از اجزای جداسازی HLO از طریق مکانیزم ساخت استاندارد OpenXLA/TensorFlow (Bazel) یا از طریق توزیع دودویی استفاده کنید.

تنظیمات منبع و Bazel

هنگام ساخت کامپایلر و پشته ابزارها از منبع، کتابخانه‌های زیر را اضافه یا به آنها وابسته کنید:

  • هدف API: //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_api
  • هدف CLI: //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_test

برای ساخت ابزار مستقل CLI:

bazel build -c opt //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_test

رابط خط فرمان (CLI)

رابط خط فرمان hlo_isolation_test به شما امکان می‌دهد تا عدم تطابق عددی و بررسی‌های پایداری را مستقیماً از یک ترمینال در برابر ماژول‌های کامپایل‌شده HLO جدا و اجرا کنید. این رابط نتایج اجرا را در محیط‌های مختلف (مثلاً TPU در مقابل TPU/CPU/Interpreter غیرفعال) مقایسه می‌کند.

مرجع پرچم

رابط خط فرمان (CLI) از پرچم‌های زیر پشتیبانی می‌کند:

  • --hlo_file : مسیر فایل ورودی .hlo یا .pbtxt برای بارگذاری. می‌تواند به صورت متنی یا پروتوتایپ باشد (الزامی).
  • --test_platform : پلتفرم هدف برای اجرای تست اصلی (مثلاً cpu ، gpu ، tpu ). پیش‌فرض روی cpu است.
  • --reference_platform : پلتفرم مرجع برای مقایسه پایه (مثلاً interpreter ). اگر خالی باشد، مقایسه مرجع غیرفعال است.
  • --filter_by_name : عبارت منظم برای مطابقت با نام ماژول. فقط ماژول‌های منطبق اجرا می‌شوند. مقدار پیش‌فرض .* است.
  • --skip_by_name : عبارت منظم برای مطابقت با نام ماژول. ماژول‌های منطبق نادیده گرفته می‌شوند.
  • --filter_by_opcode : عبارت منظم برای مطابقت با کدهای عملیاتی دستورالعمل. فقط ماژول‌هایی که حداقل شامل یک کد عملیاتی منطبق باشند اجرا می‌شوند. مقدار پیش‌فرض .* است.
  • --skip_by_opcode : عبارت منظم برای مطابقت با کدهای عملیاتی دستورالعمل. ماژول‌های حاوی هرگونه کد عملیاتی منطبق، نادیده گرفته می‌شوند.
  • --abs_error_bound : کران خطای مطلق برای مقایسه استفاده می‌شود. مقدار پیش‌فرض 0.01 است.
  • --rel_error_bound : کران خطای نسبی که برای مقایسه استفاده می‌شود. مقدار پیش‌فرض 0.1 است.
  • --run_hlo_passes : پرچم بولی برای تعیین اینکه آیا باید گذرهای استاندارد HLO را روی زیرماژول‌ها اجرا کرد یا خیر. مقدار پیش‌فرض false است.
  • --shard_index : اندیس شارد مخصوص برای اجرا (مبتنی بر صفر). مقدار پیش‌فرض -1 (غیرفعال) است.
  • --num_shards : تعداد کل Shardها. مقدار پیش‌فرض 1 است.

فراخوانی پایه 

./hlo_isolation_test \
  --hlo_file=/path/to/failing_fusion.hlo \
  --test_platform=gpu \
  --reference_platform=interpreter

نتایج و مصنوعات جمع‌آوری‌شده

وقتی یک زیرماژول در طول آزمایش جداسازی با عدم تطابق عددی یا سایر خرابی‌ها مواجه می‌شود، این ابزار به‌طور خودکار مصنوعات اشکال‌زدایی را برای بررسی عمیق‌تر روی دیسک سریالی می‌کند.

محتویات تخلیه

در صورت عدم تطابق عددی، این ابزار مصنوعات اشکال‌زدایی زیر را می‌نویسد:

  1. متن زیرماژول HLO ناموفق ( failed-module-<module_name>.txt ).
  2. خروجی مورد انتظار به صورت تحت‌اللفظی ( failed-<module_name>-expected.txt )
  3. خروجی واقعی به صورت تحت‌اللفظی ( failed-<module_name>-actual.txt ) است.
  4. خلاصه عناصر نامتناسب ( failed-<module_name>-mismatches.txt ).

محل هدف تخلیه

  • محیط تست: اگر از طریق bazel test یا محیطی که متغیر محیطی TEST_UNDECLARED_OUTPUTS_DIR را تعریف می‌کند، اجرا شود، نتایج مستقیماً در آن دایرکتوری با نام فایل‌های دقیق ذکر شده در بالا قرار می‌گیرند (مثلاً failed-<module_name>-expected.txt ).
  • اجرای استاندارد/دستی: وقتی به صورت دستی از طریق خط فرمان اجرا می‌شود، مصنوعات در دایرکتوری موقت سیستم عامل (مثلاً /tmp ) نوشته می‌شوند و دقیقاً همان قراردادهای نامگذاری یکپارچه فایل (مثلاً /tmp/failed-<module_name>-expected.txt ) را حفظ می‌کنند.

ادغام C++ و API

برای توسعه‌دهندگانی که در حال ساخت کامپایلرهای سفارشی، سکوهای تست یا خطوط لوله خودکار هستند، رابط برنامه‌نویسی کاربردی C++ راهی مستقیم برای ادغام تست‌های ایزوله فراهم می‌کند.

استفاده مستقیم از API

API اصلی، رابط‌های کاربردی برای اجرای ماژول‌ها و دریافت گزارش‌های ساختاریافته ارائه می‌دهد:

#include "third_party/tensorflow/compiler/xla/tools/hlo_isolation/hlo_isolation_api.h"

xla::hlo_isolation::PipelineIsolationOptions options;
options.module_options.abs_error_bound = 0.01;
options.module_options.rel_error_bound = 0.1;
// Filter specific opcodes programmatically
options.filter_by_opcode = "exponential";

absl::StatusOr<std::vector<xla::HloIsolationTestResult>> results =
    xla::hlo_isolation::RunIsolationPipeline(
        input_hlo_module,
        &my_test_runner,
        &my_reference_runner,
        options);

استفاده از مخلوط تست

هنگام نوشتن مجموعه‌های تست GoogleTest C++، می‌توانید از HloIsolationTestMixin برای مدیریت داخلی assertion ارث‌بری کنید. کلاس پایه باید هم یک اجراکننده تست و هم یک اجراکننده مرجع (reference runner) ارائه دهد (مثلاً از طریق HloPjRtInterpreterReferenceMixin ):

#include "third_party/tensorflow/compiler/xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
#include "third_party/tensorflow/compiler/xla/tools/hlo_isolation/hlo_isolation_test_base.h"

class MyCustomPassIsolationTest : public xla::hlo_isolation::HloIsolationTestMixin<
    xla::HloPjRtInterpreterReferenceMixin<xla::HloPjRtTestBase>> {};

TEST_F(MyCustomPassIsolationTest, ChecksMyFusionSanity) {
  RunAndVerifyIsolationTest(my_failing_module);
}

اجرای تکه تکه شده (K8s/Slurm)

برای ماژول‌های بزرگ یا ماتریس‌های تست سنگین، می‌توانید با استفاده از پرچم‌های شاردینگ، اجرا را بین خوشه‌های چند دستگاهی (مانند Google Kubernetes Engine یا Slurm) تقسیم‌بندی کنید. هر shard_index به طور قطعی یک زیرمجموعه جداگانه از زیرماژول‌های تجزیه‌شده را اجرا می‌کند. این امر امکان تأیید توزیع‌شده قابل تکرار و اجرای مجدد هدفمند پارتیشن‌های ناموفق را فراهم می‌کند.

مثال: کار Kubernetes

هر بخش آزمایشی به عنوان یک غلاف Kubernetes جداگانه با استفاده از completionMode: Indexed اجرا می‌شود. JOB_COMPLETION_INDEX مستقیماً به فلگ --shard_index در CLI ارسال می‌شود.

apiVersion: batch/v1
kind: Job
metadata:
  name: hlo-isolation-job
spec:
  completions: 50
  parallelism: 50
  completionMode: Indexed
  template:
    spec:
      containers:
      - name: test-runner
        image: gcr.io/my-project/hlo-isolation-tools:latest
        command: ["/bin/sh", "-c"]
        args:
        - |
          ./hlo_isolation_test \
            --num_shards=50 \
            --shard_index=$JOB_COMPLETION_INDEX \
            --hlo_file=/data/path/to/hlo.hlo
        volumeMounts:
        - name: hlo-data-volume
          mountPath: /data
      volumes:
      - name: hlo-data-volume
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: my-xla-debug-bucket

قابلیت‌های کلیدی

  • قابلیت حمل: بسته‌های تست داخلی را از API مستقل جدا می‌کند و اشکال‌زدایی عدم تطابق‌های HLO را به صورت محلی آسان می‌سازد.
  • جزئیات: فیلترینگ کد عملیاتی و نام جزئیات، حلقه اشکال‌زدایی را هنگام تعامل با انبوه داده‌های HLO بهبود می‌بخشد.
  • توسعه‌پذیری: فراخوانی‌های اجرای runner سفارشی و تزریق‌کننده‌های داده ( make_fake_arguments_fn ) امکان سفارشی‌سازی کامل برای گردش‌های کاری پیشرفته‌ی تأیید را فراهم می‌کنند.