یک مدل JAX آموزش دیده را صادر و ارائه دهید

۱. مقدمه

در این آزمایشگاه کد، شما از یک نقطه بازرسی آموزش‌دیده شروع می‌کنید و خط لوله استنتاج JAX را از ابتدا تا انتها، از یک مسیر رو به جلو کامپایل شده توسط JIT تا کامپایل پیش از زمان (AOT)، خروجی بومی JAX با jax.export و یک مدل ذخیره شده TensorFlow با jax2tf ، طی می‌کنید. هر مسیر اندازه‌گیری می‌شود و هر چهار مسیر برای تولید پیش‌بینی‌های یکسان در نظر گرفته شده‌اند.

کاری که انجام خواهید داد

• ترانسفورماتور codelab 7 را بازسازی کنید و وزن‌های آموزش‌دیده آن را با Orbax بارگذاری کنید.
• مرحله‌ی رو به جلو را در jax.jit قرار دهید و تأخیر اولین فراخوانی را در مقابل تأخیر فراخوانی ذخیره‌شده اندازه‌گیری کنید.
• شروع سرد با کامپایل AOT را حذف کنید ( lower() then compile() ) و StableHLO IR را بخوانید
• اندازه‌گیری توان عملیاتی رو به جلو بر حسب توکن بر ثانیه در چهار اندازه دسته‌ای
• یک مصنوع قابل حمل را با jax.export سریالی کنید، سپس آن را deserialize و فراخوانی کنید
• مدل را با jax2tf به TensorFlow SavedModel تبدیل کنید و هر چهار مسیر را با هم مقایسه کنید.

آنچه نیاز دارید

• یک پروژه Google Cloud با قابلیت پرداخت صورتحساب، و اعتبار کارگاه یا رزرو شامل استفاده از GPU
• سهمیه حداقل ۲ پردازنده گرافیکی NVIDIA L4 در منطقه انتخابی شما ( نحوه بررسی سهمیه پردازنده گرافیکی )
• تکمیل آزمایشگاه‌های کد ۱ تا ۷ یا یک محیط معادل JAX GPU مجهز به CUDA
• اجرای کامل codelab 7، که این آزمایشگاه، Checkpoint مربوط به Orbax آن را بارگذاری می‌کند.

زمان تخمینی برای تکمیل: ۶۰ دقیقه .

چهار مسیر از آموزش تا خدمت

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

مسیر	قالب	هدف خدمت	چه زمانی استفاده شود
jax.jit	فایل اجرایی ذخیره شده در حین فرآیند	سرور پایتون (FastAPI، Flask)	ساده‌ترین مسیر سرویس‌دهی JAX با تأخیر کم
AOT compile	فایل اجرایی از پیش کامپایل شده در حین پردازش	راه‌اندازی/گرم کردن سرور پایتون	جلوگیری از تأخیر در کامپایل اولین درخواست
jax.export	خروجی JAX سریالی شده با متادیتای StableHLO +	زمان اجرای JAX سازگار برای پلتفرم(های) صادر شده	مصنوع قابل حمل بومی JAX
jax2tf	TF مدل ذخیره شده	خدمات TF، TFX	اکوسیستم تنسورفلو

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

۲. قبل از شروع

پروژه خود را انتخاب کنید

در کنسول گوگل کلود ، یک پروژه با قابلیت پرداخت فعال انتخاب یا ایجاد کنید.

پوسته ابری را باز کنید

برای شروع یک جلسه Cloud Shell ، روی Activate Cloud Shell (آیکون ترمینال در سمت راست بالای کنسول) کلیک کنید، سپس آن را به پروژه خود هدایت کنید:

gcloud config set project <YOUR_PROJECT_ID>

فراهم کردن محیط GPU

دستور زیر را در Cloud Shell اجرا کنید. Codelab 1 هر دستور را با جزئیات توضیح می‌دهد، از جمله سهمیه GPU و الزامات منطقه.

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform

cp terraform.tfvars.example terraform.tfvars

فایل terraform.tfvars را ویرایش کنید و project_id = " " قرار دهید. project_id = " " ، سپس کلاستر را آماده کرده و JupyterLab را مستقر کنید:

terraform init
terraform apply
$(terraform output -raw get_credentials_command)

cd ..
kubectl apply -f deploy/jupyter.yaml

terraform apply حدود ۱۲ دقیقه طول می‌کشد. پس از اتمام، منتظر Pod و LoadBalancer باشید، سپس توکن یکبار مصرف JupyterLab را از لاگ Pod بخوانید:

kubectl get pod jax-jupyter -w         # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w     # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1

http:// :8884 را باز کنید http:// :8884 وارد کنید، توکن را جایگذاری کنید و یک دفترچه یادداشت پایتون ۳ جدید در /workspace ایجاد کنید. هر بلوک کد در این codelab به سلولی از آن دفترچه یادداشت می‌رود.

آنچه این codelab نیاز دارد را نصب کنید

!pip install --quiet flax orbax-checkpoint tensorflow

flax و orbax-checkpoint معمولاً در کانتینر NVIDIA JAX ارائه می‌شوند، بنابراین این دو معمولاً نیازی به عملیات ندارند. اگر TensorFlow از قبل وجود نداشته باشد، pip آن را دریافت می‌کند.

GPU را تنظیم و تأیید کنید

JAX، Flax NNX و Orbax را وارد کنید، تأیید کنید که حداقل یک GPU قابل مشاهده است و دو کمکی را که بقیه codelab برای مسدود کردن نتایج و رندر جداول نتیجه استفاده می‌کنند، تعریف کنید.

import os
os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")

import html
import pathlib
import time
import warnings

from IPython.display import HTML, display
import numpy as np

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", message=".*ml_dtypes.*")
warnings.filterwarnings("ignore", message=".*JAX_PLATFORMS.*")

import jax
import jax.numpy as jnp
from flax import nnx
import orbax.checkpoint as ocp


devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]

print(f"JAX version:     {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"GPU devices:     {gpu_devices}")

assert len(gpu_devices) >= 1, (
    f"This lesson needs at least 1 GPU. Found {len(gpu_devices)}. "
    f"Available devices: {devices}"
)


def block_tree(tree):
    return jax.block_until_ready(tree)


def show_table(headers, rows, title=None, aligns=None):
    aligns = aligns or ["left"] * len(headers)
    parts = ["<div style='font-family: system-ui; max-width: 980px;'>"]
    if title:
        parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    parts.append("<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>")
    parts.append("<thead><tr>")
    for h, a in zip(headers, aligns):
        parts.append(
            f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>"
            f"{html.escape(str(h))}</th>"
        )
    parts.append("</tr></thead><tbody>")
    for row in rows:
        parts.append("<tr>")
        for cell, a in zip(row, aligns):
            parts.append(
                f"<td style='text-align:{a}; border-bottom:1px solid #eef1f4; padding:6px;'>"
                f"{html.escape(str(cell))}</td>"
            )
        parts.append("</tr>")
    parts.append("</tbody></table></div>")
    display(HTML("\n".join(parts)))

شما باید یک نسخه JAX، gpu به عنوان backend پیش‌فرض و لیستی با حداقل یک دستگاه CUDA را ببینید.

تأیید کنید که یک ایست بازرسی آموزش دیده دارید

این codelab، Checkpoint مربوط به Orbax را بارگذاری می‌کند که یک ترانسفورماتور را از ابتدا تا انتها با Flax NNX آموزش می‌دهد و Orbax codelab آن را در /tmp/jax-course/l7-checkpoints/trained می‌نویسد، بنابراین قبل از ساخت هر چیزی، بررسی کنید که دایرکتوری هنوز وجود داشته باشد.

import pathlib
ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")
assert (ckpt_dir / "trained").exists(), (
    f"No checkpoint at {ckpt_dir / 'trained'}. Run codelab 7 first."
)
print(f"Found checkpoint: {ckpt_dir / 'trained'}")

باید ببینید که مسیر نقطه بازرسی برای شما چاپ می‌شود.

۳. مدل را بازسازی کنید و Checkpoint را بارگذاری کنید

یک Checkpoint در Orbax مقادیر پارامترها را ذخیره می‌کند، نه کلاس مدلی که آنها را تولید کرده است. برای بازیابی آن، ابتدا به همان معماری نیاز دارید، بنابراین این مرحله Codelab 7 TinyTransformer دوباره تعریف می‌کند و سپس آن را با وزن‌های ذخیره شده پر می‌کند.

معماری را از نو تعریف کنید

این مدل دارای یک جاسازی توکن و موقعیت، چهار بلوک تبدیل پیش هنجار با توجه علّی، یک هنجار لایه نهایی و یک سر مدل زبانی است.

VOCAB_SIZE = 256
D_MODEL = 256
NUM_HEADS = 4
FFN_DIM = 1024
NUM_LAYERS = 4
MAX_SEQ_LEN = 256


def causal_sdpa(query, key, value, **_):
    return jax.nn.dot_product_attention(query, key, value, is_causal=True)


class TransformerBlock(nnx.Module):
    def __init__(self, d_model: int, num_heads: int, ffn_dim: int, rngs: nnx.Rngs):
        self.ln1 = nnx.LayerNorm(d_model, rngs=rngs)
        self.attn = nnx.MultiHeadAttention(
            num_heads=num_heads,
            in_features=d_model,
            decode=False,
            attention_fn=causal_sdpa,
            rngs=rngs,
        )
        self.ln2 = nnx.LayerNorm(d_model, rngs=rngs)
        self.fc_up = nnx.Linear(d_model, ffn_dim, rngs=rngs)
        self.fc_down = nnx.Linear(ffn_dim, d_model, rngs=rngs)

    def __call__(self, x):
        x = x + self.attn(self.ln1(x))
        h = jax.nn.gelu(self.fc_up(self.ln2(x)))
        x = x + self.fc_down(h)
        return x


class TinyTransformer(nnx.Module):
    def __init__(self, vocab_size: int, d_model: int, num_heads: int,
                 ffn_dim: int, num_layers: int, max_seq_len: int,
                 rngs: nnx.Rngs):
        self.token_embed = nnx.Embed(vocab_size, d_model, rngs=rngs)
        self.pos_embed = nnx.Embed(max_seq_len, d_model, rngs=rngs)
        self.blocks = nnx.List([
            TransformerBlock(d_model, num_heads, ffn_dim, rngs=rngs)
            for _ in range(num_layers)
        ])
        self.final_norm = nnx.LayerNorm(d_model, rngs=rngs)
        self.lm_head = nnx.Linear(d_model, vocab_size, use_bias=False, rngs=rngs)

    def __call__(self, tokens):
        B, T = tokens.shape
        x = self.token_embed(tokens) + self.pos_embed(jnp.arange(T))
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        return self.lm_head(x)


param_count = sum(x.size for x in jax.tree.leaves(nnx.state(TinyTransformer(
    VOCAB_SIZE, D_MODEL, NUM_HEADS, FFN_DIM, NUM_LAYERS, MAX_SEQ_LEN,
    rngs=nnx.Rngs(0),
), nnx.Param)))
print(f"TinyTransformer: {param_count:,} parameters")

شما باید تعداد پارامترها را برای یک مدل تازه مقداردهی شده ببینید. وزن‌ها در این مرحله تصادفی هستند زیرا معماری در اینجا اهمیت دارد.

وزنه‌های تمرین‌شده را بازیابی کنید

Orbax مدل را به یک ساختار هدف (target structure) بازیابی می‌کند. شما آن هدف را با گرفتن وضعیت مدل جدید و جایگزینی هر آرایه با یک jax.ShapeDtypeStruct می‌سازید، که یک شکل و نوع داده (dtype) را بدون تخصیص حافظه توصیف می‌کند. سپس Orbax نقطه بررسی را به آرایه‌هایی دقیقاً از همان شکل‌ها می‌خواند.

ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")

model = TinyTransformer(
    VOCAB_SIZE, D_MODEL, NUM_HEADS, FFN_DIM, NUM_LAYERS, MAX_SEQ_LEN,
    rngs=nnx.Rngs(0),
)
model_params = nnx.state(model, nnx.Param)
abstract_params = jax.tree.map(
    lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), model_params
)

checkpointer = ocp.StandardCheckpointer()
restored_params = checkpointer.restore(ckpt_dir / "trained", abstract_params)

# Move to a single GPU
single_device = gpu_devices[0]
restored_params = jax.device_put(restored_params, single_device)
nnx.update(model, restored_params)

print(f"\u2705 Checkpoint loaded from {ckpt_dir / 'trained'}")
print(f"Parameters: {sum(x.size for x in jax.tree.leaves(restored_params)):,}")

فراخوانی jax.device_put مهم است. Serving از یک GPU استفاده می‌کند، بنابراین پارامترها قبل از ورود به مدل به gpu_devices[0] منتقل می‌شوند.

شما باید مسیر ایست بازرسی تأیید شده و تعداد پارامترها را که با تعداد چاپ شده توسط سلول قبلی مطابقت دارد، ببینید.

۴. مسیر ۱: با jax.jit سرو کنید

ساده‌ترین مسیر سرویس‌دهی، قرار دادن مسیر رو به جلو در jax.jit است. اولین فراخوانی، کامپایل را آغاز می‌کند و هر فراخوانی پس از آن، فایل اجرایی ذخیره شده را دوباره استفاده می‌کند.

برای ارائه، وزن‌ها را در یک بستار قرار دهید. nnx.split مدل را به یک تعریف گراف استاتیک و یک حالت دینامیک تقسیم می‌کند و تابع jitted هر دو را دریافت می‌کند و فقط tokens به عنوان آرگومان دریافت می‌کند. این باعث می‌شود تابع کامپایل شده مستقل باشد.

graphdef, model_state = nnx.split(model)

@jax.jit
def predict_jit(tokens):
    m = nnx.merge(graphdef, model_state)
    return m(tokens)

dummy_input = jnp.zeros((1, MAX_SEQ_LEN), dtype=jnp.int32)

start = time.perf_counter()
logits = block_tree(predict_jit(dummy_input))
first_call_ms = (time.perf_counter() - start) * 1000

times = []
for _ in range(100):
    start = time.perf_counter()
    logits = block_tree(predict_jit(dummy_input))
    times.append((time.perf_counter() - start) * 1000)

avg_ms = np.mean(times)

show_table(
    ["", "Latency (ms)"],
    [
        ("First call (compile + execute)", f"{first_call_ms:,.1f}"),
        ("Subsequent calls (avg of 100)", f"{avg_ms:.2f}"),
        ("Speedup", f"{first_call_ms / avg_ms:.0f}\u00d7"),
    ],
    title="JIT inference latency",
    aligns=["left", "right"],
)

print(f"\nOutput shape: {logits.shape} (batch=1, seq={MAX_SEQ_LEN}, vocab={VOCAB_SIZE})")

شما باید یک جدول تأخیر استنتاج JIT را ببینید که در آن اولین فراخوانی بسیار کندتر از میانگین ۱۰۰ فراخوانی بعدی است، و سپس Output shape: (1, 256, 256) .

۵. مسیر ۲: حذف شروع سرد با کامپایل AOT

jax.jit در اولین فراخوانی کامپایل می‌شود، که در طول توسعه مشکلی ندارد و در زمان اجرا با مشکل مواجه می‌شود. کامپایل AOT آن مرحله را به مراحل جداگانه تقسیم می‌کند:

1. شما یک تابع معمولی پایتون/JAX می‌نویسید، مانند یک تابع model forward pass
2. JAX تابع را برای شکل‌ها و نوع‌های ورودی خاص ردیابی می‌کند و آن را با استفاده از lower() به IR کامپایلر تبدیل می‌کند.
3. StableHLO محاسبات را در عملیات مستقل از سخت‌افزار مانند dot ، reshape و reduce توصیف می‌کند.
4. XLA، StableHLO را بهینه می‌کند و با استفاده از compile() یک فایل اجرایی مخصوص دستگاه برای GPU، TPU یا CPU تولید می‌کند.

نمایش میانی StableHLO (IR) نمایش سطح کامپایلر یک محاسبه JAX پس از تبدیل کد پایتون به یک برنامه قابل حمل و مستقل از سخت‌افزار است. این نمایش، عملیاتی مانند ضرب ماتریس، تغییر شکل، کاهش و جریان کنترل را به شکلی توصیف می‌کند که XLA می‌تواند برای بک‌اندهای مختلف، از جمله GPUها، TPUها و CPUها، کامپایل کند.

شما یک بار، یا در هنگام راه‌اندازی، یا به صورت آفلاین برای یک شکل ورودی ثابت، کامپایل می‌کنید و سپس بارها و بارها بدون پرداخت سربار کامپایل فراخوانی اول، اجرا می‌کنید.

# Bake weights into a closure
def predict_closed(tokens):
    m = nnx.merge(graphdef, model_state)
    return m(tokens)

# Stage 1: Lower
abstract_tokens = jax.ShapeDtypeStruct((1, MAX_SEQ_LEN), jnp.int32)

lowered = predict_closed.lower(abstract_tokens)
print(f"Lowered to StableHLO ({len(lowered.as_text()):,} chars)")

# Stage 2: Compile
compiled = lowered.compile()
print(f"Compiled for: {jax.default_backend()}")

# Execute
start = time.perf_counter()
logits_aot = block_tree(compiled(dummy_input))
aot_first_ms = (time.perf_counter() - start) * 1000

times_aot = []
for _ in range(100):
    start = time.perf_counter()
    logits_aot = block_tree(compiled(dummy_input))
    times_aot.append((time.perf_counter() - start) * 1000)

avg_aot_ms = np.mean(times_aot)
max_diff_jit_aot = float(jnp.max(jnp.abs(logits - logits_aot)))

show_table(
    ["", "Latency (ms)"],
    [
        ("AOT first execution (no compile)", f"{aot_first_ms:.2f}"),
        ("AOT subsequent (avg of 100)", f"{avg_aot_ms:.2f}"),
        ("JIT first call (from above)", f"{first_call_ms:,.1f}"),
        ("Max |JIT − AOT|", f"{max_diff_jit_aot:.2e}"),
    ],
    title="AOT vs JIT latency",
    aligns=["left", "right"],
)

توجه داشته باشید که lower() یک jax.ShapeDtypeStruct می‌گیرد، نه داده‌های واقعی. شما هرگز به آرایه ورودی برای کامپایل نیاز ندارید - فقط شکل و نوع داده آن.

شما باید جدول تأخیر AOT در مقابل JIT را ببینید که در آن اولین اجرای AOT به عدد حالت پایدار AOT نزدیک است تا به اولین فراخوانی JIT و مقدار Max |JIT − AOT| که عملاً صفر است.

StableHLO IR را بررسی کنید

lowered.as_text() برنامه StableHLO را که روی دستگاه اجرا خواهد شد، نشان می‌دهد. این همان نمایش میانی است که XLA در GPUها، TPUها و CPUها استفاده می‌کند و برای اشکال‌زدایی، تحلیل عملکرد و درک آنچه کامپایلر واقعاً می‌بیند، مفید است.

متن StableHLO می‌تواند بسیار بزرگ باشد و یک خط ممکن است شامل یک ثابت یا ویژگی طولانی باشد. برای جلوگیری از محدودیت نرخ داده IOPub ژوپیتر، کد بعدی کل IR را روی دیسک ذخیره می‌کند و فقط یک پیش‌نمایش محدود چاپ می‌کند.

hlo_text = lowered.as_text()

hlo_path = pathlib.Path("/tmp/jax-course/l8-stablehlo.mlir")
hlo_path.parent.mkdir(parents=True, exist_ok=True)
hlo_path.write_text(hlo_text)

MAX_LINES = 20
MAX_CHARS_PER_LINE = 160

lines = hlo_text.splitlines()
preview_lines = []
for line in lines[:MAX_LINES]:
    if len(line) > MAX_CHARS_PER_LINE:
        preview_lines.append(line[:MAX_CHARS_PER_LINE] + " ... [line truncated]")
    else:
        preview_lines.append(line)

print(f"StableHLO program: {len(lines):,} lines, {len(hlo_text):,} chars")
print(f"Full StableHLO saved to: {hlo_path}")
print("=" * 60)
print("\n".join(preview_lines))
print(
    f"\n... ({max(len(lines) - MAX_LINES, 0):,} more lines; "
    "long lines are truncated in this preview)"
)

شما باید تعداد خطوط و کاراکترهای برنامه، مسیر فایل ذخیره شده با پسوند .mlir و 20 خط اول IR و به دنبال آن تعداد خطوط جا افتاده را ببینید.

۶. اندازه‌گیری توان عملیاتی استنتاج دسته‌ای

پردازنده‌های گرافیکی (GPU) هنگام پردازش همزمان چندین ورودی، بیشترین کارایی را دارند. این مرحله، میزان مقیاس‌پذیری توان عملیاتی رو به جلو را با اندازه دسته اندازه‌گیری می‌کند. تابع مدل تغییر نمی‌کند، اما هر شکل ورودی جدید به فایل اجرایی کاهش یافته و کامپایل شده مخصوص به خود نیاز دارد، بنابراین حلقه به ازای هر اندازه دسته، یک بار کامپایل می‌شود.

batch_sizes = [1, 4, 16, 64]
results = []

for bs in batch_sizes:
    tokens_batch = jnp.zeros((bs, MAX_SEQ_LEN), dtype=jnp.int32)

    # Compile for this batch size
    lowered_bs = predict_closed.lower(
        jax.ShapeDtypeStruct((bs, MAX_SEQ_LEN), jnp.int32),
    )
    compiled_bs = lowered_bs.compile()

    # Warmup
    block_tree(compiled_bs(tokens_batch))

    # Measure
    times_bs = []
    for _ in range(50):
        start = time.perf_counter()
        block_tree(compiled_bs(tokens_batch))
        times_bs.append((time.perf_counter() - start) * 1000)

    avg_bs = np.mean(times_bs)
    tokens_per_sec = (bs * MAX_SEQ_LEN) / (avg_bs / 1000)
    results.append((bs, f"{avg_bs:.2f}", f"{tokens_per_sec:,.0f}"))

show_table(
    ["Batch size", "Latency (ms)", "Tokens/sec"],
    results,
    title="Batched inference throughput",
    aligns=["right", "right", "right"],
)

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

۷. مسیر ۳: خروجی گرفتن از یک مصنوع قابل حمل با jax.export

jax.export یک تابع JAX جیت‌شده را به یک شیء Exported حاوی StableHLO به همراه فراداده‌های مورد نیاز برای فراخوانی آن از یک فرآیند JAX دیگر صادر می‌کند. بایت‌های سریالیزه شده می‌توانند:

• روی دیسک ذخیره شده و در فرآیند دیگری بارگذاری می‌شود
• بدون کد منبع مدل اصلی پایتون فراخوانی می‌شود
• به طور پیش‌فرض برای پلتفرم فعلی یا برای پلتفرم‌های مشخص با آرگومان platforms=[...] صادر می‌شود.

این مسیر استقرار بومی JAX است - هیچ وابستگی به TensorFlow لازم نیست.

from jax import export

# Export the closure-based function
exported = export.export(predict_closed)(
    jax.ShapeDtypeStruct((1, MAX_SEQ_LEN), jnp.int32),
)

print(f"Exported function: {exported.fun_name}")
print(f"Input shapes:  {exported.in_avals}")
print(f"Output shapes: {exported.out_avals}")
print(f"Exported platforms: {exported.platforms}")

# Serialize to bytes
blob = exported.serialize()
export_path = pathlib.Path("/tmp/jax-course/exports")
export_path.mkdir(parents=True, exist_ok=True)

export_file = export_path / "tiny_transformer_jax_export.bin"
export_file.write_bytes(blob)
print()
print(f"Serialized to {export_file} ({len(blob):,} bytes, {len(blob) / 1024:.0f} KB)")

# Deserialize and call
rehydrated = export.deserialize(export_file.read_bytes())

test_input = jnp.zeros((1, MAX_SEQ_LEN), dtype=jnp.int32)
logits_exported = block_tree(rehydrated.call(test_input))
print()
print(f"✅ Deserialized call succeeded — output shape: {logits_exported.shape}")

# Verify outputs match
diff_export = float(jnp.max(jnp.abs(logits_aot - logits_exported)))
print(f"Max difference from AOT: {diff_export:.2e}")

شیء rehydrated هرگز TinyTransformer ندید. این شیء از بایت‌های روی دیسک بازسازی شده و هنوز همان logits را برمی‌گرداند، که کل نکته یک مصنوع قابل حمل است.

شما باید نام تابع خروجی، مقادیر ورودی و خروجی، لیست پلتفرم خروجی، تعداد بایت‌های blob سریال‌سازی شده، شکل خروجی از فراخوانی deserialized و حداکثر اختلاف از logit های AOT که عملاً صفر است را ببینید.

۸. مسیر ۴: تبدیل به یک TensorFlow SavedModel با jax2tf

اگر زیرساخت سرویس‌دهی شما از TensorFlow مانند TF Serving یا TFX pipelines استفاده می‌کند، می‌توانید تابع JAX را به یک TF SavedModel تبدیل کنید. jax2tf هنوز در jax.experimental قرار دارد، اما مسیر استاندارد تعامل JAX با TensorFlow است.

سریال‌سازی بومی در نسخه‌های فعلی JAX پیش‌فرض است، بنابراین جاسازی‌های تبدیل، StableHLO را در نمودار TensorFlow بدون عبور از native_serialization=True کاهش دادند.

یک جزئیات پلتفرم در اینجا اهمیت دارد. این codelab، JAX را روی CUDA اجرا می‌کند، اما زمان اجرای TensorFlow در کانتینر JAX ممکن است SavedModel را روی CPU اجرا کند. برای جلوگیری از فراخوانی ماژول صادر شده توسط CUDA توسط TensorFlow روی CPU، کد زیر ماژول jax2tf را برای ("cpu",) صادر می‌کند. اگر با TensorFlow روی GPU کار می‌کنید، برای ("cuda",) صادر کنید و از یک زمان اجرای TensorFlow با پشتیبانی از GPU/XLA استفاده کنید.

from jax.experimental import jax2tf
import tensorflow as tf
import shutil

# Capture model_state as a closure
def predict_for_tf(tokens):
    m = nnx.merge(graphdef, model_state)
    return m(tokens)

TF_EXPORT_PLATFORMS = ("cpu",)
tf_predict = jax2tf.convert(
    predict_for_tf,
    native_serialization_platforms=TF_EXPORT_PLATFORMS,
)

# Wrap in a tf.Module for SavedModel export
module = tf.Module()
module.predict = tf.function(
    tf_predict,
    input_signature=[tf.TensorSpec(shape=(1, MAX_SEQ_LEN), dtype=tf.int32)],
    autograph=False,
)

# TF Serving expects a versioned model directory
savedmodel_base_dir = pathlib.Path("/tmp/jax-course/exports") / "tiny_transformer_savedmodel"
savedmodel_dir = savedmodel_base_dir / "1"
if savedmodel_base_dir.exists():
    shutil.rmtree(savedmodel_base_dir)

tf.saved_model.save(module, str(savedmodel_dir))
print(f"✅ SavedModel saved to {savedmodel_dir}")
print(f"Exported for TensorFlow platform(s): {TF_EXPORT_PLATFORMS}")

# Verify the SavedModel path produces the same logits as the JAX path
with tf.device("/CPU:0"):
    tf_logits = module.predict(tf.zeros((1, MAX_SEQ_LEN), dtype=tf.int32))
diff_tf = np.max(np.abs(np.asarray(tf_logits) - np.asarray(logits_aot)))
print(f"Max difference from AOT: {diff_tf:.2e}")

# List saved files
for dirpath, _, filenames in os.walk(savedmodel_base_dir):
    for f in filenames:
        full = os.path.join(dirpath, f)
        size = os.path.getsize(full)
        print(f"  {os.path.relpath(full, savedmodel_base_dir):44s} {size:>10,} bytes")
print()
print(
    f"To serve: docker run -p 8501:8501 "
    f"--mount type=bind,source={savedmodel_base_dir},target=/models/transformer "
    "-e MODEL_NAME=transformer tensorflow/serving "
    "--xla_cpu_compilation_enabled=true"
)

شما باید مسیر دایرکتوری SavedModel، تاپل پلتفرم صادرات، حداکثر اختلاف از logits AOT که عملاً صفر است، و لیستی از فایل‌های نوشته شده در دایرکتوری نسخه‌بندی شده را ببینید.

آخرین چیزی که کد چاپ می‌کند، دستور docker run ... tensorflow/serving است. این دستور توضیحی است و نشان می‌دهد که چگونه TensorFlow Serving را به دایرکتوری نسخه‌بندی‌شده‌ای که ایجاد کرده‌اید، ارجاع می‌دهید. شما آن را در این آزمایشگاه کد اجرا نمی‌کنید.

۹. چهار مسیر سرو را با هم مقایسه کنید

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

show_table(
    ["Path", "Format", "Dependencies", "Serving target", "Portable"],
    [
        ("jax.jit", "Cached in-process executable", "JAX", "Python server", "No"),
        ("AOT compile", "Compiled in-process executable", "JAX", "Python server startup / warmup", "No"),
        ("jax.export", "Serialized JAX export", "JAX runtime", "Compatible runtime for exported platform(s)", "Yes"),
        ("jax2tf", "TF SavedModel", "TensorFlow", "TF Serving", "Yes"),
    ],
    title="Serving path comparison",
)

# Show file sizes
sizes = []
export_size = os.path.getsize(export_file)
sizes.append(("jax.export", f"{export_size:,} bytes", f"{export_size / 1024:.0f} KB"))

sm_size = sum(
    os.path.getsize(os.path.join(dirpath, f))
    for dirpath, _, filenames in os.walk(savedmodel_base_dir)
    for f in filenames
)
sizes.append(("jax2tf SavedModel", f"{sm_size:,} bytes", f"{sm_size / 1024:.0f} KB"))

show_table(
    ["Export", "Size (bytes)", "Size (KB)"],
    sizes,
    title="Export file sizes",
    aligns=["left", "right", "right"],
)

شما باید جدول مقایسه مسیر Serving و به دنبال آن جدول اندازه‌های فایل Export را ببینید که اندازه روی دیسک jax.export blob و دایرکتوری SavedModel را فهرست می‌کند.

بر اساس محدودیت‌های استقرار خود انتخاب کنید، نه بر اساس معماری مدل. اگر JAX از قبل در فرآیند ارائه شما در حال اجرا است، jax.jit به همراه یک startup warmup یا کامپایل AOT کوتاه‌ترین مسیر است. اگر مصنوع مجبور به ترک آن فرآیند باشد، jax.export شما را در داخل JAX نگه می‌دارد و jax2tf شما را به اکوسیستم TensorFlow تحویل می‌دهد.

۱۰. تمیز کردن

هر چیزی که این آزمایشگاه کد نوشته است در پوشه /tmp در Pod قرار دارد که با Pod ناپدید می‌شود. ابتدا هر چیزی را که می‌خواهید نگه دارید، کپی کنید و این دستور را از Cloud Shell اجرا کنید:

kubectl cp jax-jupyter:/tmp/jax-course/exports ./jax-course-exports

حجم کاری Jupyter، شامل LoadBalancer و درایو دائمی را حذف کنید:

kubectl delete -f deploy/jupyter.yaml

کلاستر، Node Pool، VPC و حساب کاربری سرویس را از بین ببرید:

cd terraform
terraform destroy

وقتی از شما خواسته شد، yes را تایپ کنید، سپس تأیید کنید که چیزی جا نمانده است:

gcloud container clusters list
gcloud compute instances list

هر دو باید برای این پروژه خالی باشند. اگر فقط برای این مجموعه پروژه‌ای ایجاد کرده‌اید، می‌توانید کل پروژه را از کنسول Cloud حذف کنید.

۱۱. تبریک

شما یک نقطه کنترل ترانسفورماتور آموزش‌دیده را از نوت‌بوک برداشتید و آن را به چهار مصنوع آماده برای استفاده در یک پردازنده گرافیکی NVIDIA L4 تبدیل کردید.

آنچه آموخته‌اید

• jax.jit ساده‌ترین مسیر است - wrap و call. اولین فراخوانی کامپایل می‌شود؛ فراخوانی‌های بعدی سریع هستند. برای سرویس‌های مبتنی بر پایتون (FastAPI، Flask) که JAX از قبل نصب شده است، مناسب است.
• کامپایل AOT ( lower() then compile() ) کامپایل را از اجرا جدا می‌کند. یک بار در هنگام راه‌اندازی کامپایل می‌شود، سپس بدون تأخیر کامپایل در اولین درخواست، اجرا می‌شود. lowered.as_text() StableHLO IR را برای اشکال‌زدایی و تحلیل عملکرد نشان می‌دهد.
• jax.export یک تابع jitted را به یک مصنوع بومی JAX حاوی StableHLO و فراداده‌های فراخوانی سریالی می‌کند. فایل حاصل می‌تواند توسط یک زمان اجرای سازگار با JAX بدون کد مدل اصلی بارگیری و فراخوانی شود. به طور پیش‌فرض برای پلتفرم فعلی صادر می‌شود. وقتی به یک هدف صریح نیاز دارید platforms=[...] استفاده کنید.
• jax2tf تابع را به یک TensorFlow SavedModel تبدیل می‌کند. وقتی زیرساخت سرویس‌دهی شما مبتنی بر TensorFlow است، از آن استفاده کنید. سریال‌سازی بومی (Native serialization) پیش‌فرض در نسخه‌های فعلی JAX است؛ native_serialization_platforms را طوری تنظیم کنید که با جایی که TensorFlow مدل را اجرا می‌کند، مطابقت داشته باشد.
• نحوه بازیابی یک ایست بازرسی Orbax به یک معماری بازسازی شده با استفاده از اهداف jax.ShapeDtypeStruct و نحوه انتقال پارامترها به یک دستگاه واحد با jax.device_put
• چرا forward-pass tokens/sec به عنوان توان عملیاتی تولید خودهمبسته (اتورگرسیو) در نظر گرفته نمی‌شود، و چرا هر شکل ورودی جدید به فایل اجرایی کامپایل‌شده‌ی مخصوص به خود نیاز دارد؟

خلاصه دوره

بیش از هشت آزمایشگاه، شما موارد زیر را دارید:

1. L1-L3 : راه‌اندازی JAX، یادگیری کامپایل jit و اجرای GPU پروفایل‌بندی‌شده
2. L4 : یک حلقه آموزشی از ابتدا ساخته شده است - بهینه ساز، ضرر، به روز رسانی گرادیان
3. L5 : بررسی مکانیسم‌های توجه - توجه ترکیبی ساده، SDPA، cuDNN
4. L6 : مقیاس‌پذیر برای چندین پردازنده گرافیکی با موازی‌سازی داده‌ها - مش، NamedSharding، شاردینگ موازی داده‌ها
5. L7 : همه چیز را در یک مدل زبان ترانسفورماتور ترکیب کرد - Flax NNX، Orbax، generation
6. L8 : مدل آموزش‌دیده را برای تولید آماده کرد — JIT، AOT، jax.export ، jax2tf

مراحل بعدی

اکنون گردش کار اصلی برای ساخت با JAX روی GPUها را دارید. بهترین قدم بعدی تبدیل این قطعات به پروژه‌های بزرگ‌تر و شلوغ‌تر است:

• رمزگشایی KV-cache و تولید دسته‌ای را اضافه کنید، تا بتوانید توان عملیاتی تولید واقعی را به جای توان عملیاتی رو به جلو اندازه‌گیری کنید.
• به جای واژگان اسباب‌بازی که اینجا استفاده می‌شود، روی یک توکن‌ساز و مجموعه داده واقعی آموزش دهید
• با دقت و کمیت ترکیبی آزمایش کنید و به جای فرض افزایش سرعت، آن را اندازه‌گیری کنید
• فراتر از موازی‌سازی داده‌ها، به موازی‌سازی مدل و تانسور بپردازید
• یک پشته سرو کوچک را در اطراف یکی از مصنوعات امروزی، به همراه نظارت و آزمایش‌های بارگذاری، بسازید.

از اینجا تمرکز از یادگیری ویژگی‌های JAX به سمت انتخاب‌های مهندسی تغییر می‌کند: چگونه مدل‌ها را سریع، قابل تکرار، با حافظه کارآمد، قابل اشکال‌زدایی و قابل استقرار نگه داریم.

اسناد مرجع

• کامپایل JAX AOT
• راهنمای jax.export
• یادداشت‌های تغییرات JAX برای سریال‌سازی بومی jax2tf
• اصول اولیه Flax NNX
• راهنمای ایست بازرسی Orbax
• کانتینر NVIDIA JAX روی NGC