Improve performance on Trainium2 with Logical Neuron Cores (LNC)

============================================================
  LNC Benchmark -- SigLIP on Trainium2
============================================================

============================================================
  LNC=2  (4 logical cores)
============================================================
Downloading ViT-SO400M-14-SigLIP-384 (webli)...
  Vision encoder ready (428M params)
  Compiling LNC=2 (flags: --optlevel 3 --model-type transformer --auto-cast matmult --lnc 2)...
  Saved siglip_lnc2.pt (1333 MB, 244s)

  Single-core benchmark (10 warmup, 50 timed)...
    Throughput: 22.55 img/s
    Latency:    44.35 ms

  Whole-instance benchmark (4 cores, 60s)...
    Core 0: 22.54 img/s
    Core 1: 22.54 img/s
    Core 2: 22.57 img/s
    Core 3: 22.56 img/s
    Total: 90.21 img/s

============================================================
  LNC=1  (8 logical cores)
============================================================
Downloading ViT-SO400M-14-SigLIP-384 (webli)...
  Vision encoder ready (428M params)
  Compiling LNC=1 (flags: --optlevel 3 --model-type transformer --auto-cast matmult --lnc 1)...
  Saved siglip_lnc1.pt (672 MB, 188s)

  Single-core benchmark (10 warmup, 50 timed)...
    Throughput: 26.72 img/s
    Latency:    37.42 ms

  Whole-instance benchmark (8 cores, 60s)...
    Core 0: 26.63 img/s
    Core 1: 26.61 img/s
    Core 2: 26.63 img/s
    Core 3: 26.61 img/s
    Core 4: 26.52 img/s
    Core 5: 26.49 img/s
    Core 6: 26.52 img/s
    Core 7: 26.51 img/s
    Total: 212.53 img/s

============================================================
  COMPARISON: LNC=2 (default) vs LNC=1
============================================================
  Model: ViT-SO400M-14-SigLIP-384
  Instance: trn2.3xlarge

  Metric                                LNC=2        LNC=1
  ------------------------------------------------------
  Logical cores                             4            8
  Compiled model size              1332.91 MB    671.75 MB
  Single-core throughput          22.55 img/s  26.72 img/s
  Single-core latency                44.35 ms     37.42 ms
  Whole-instance throughput       90.21 img/s 212.53 img/s
  ------------------------------------------------------

source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_13/bin/activate
pip install open_clip_torch
python3 lnc_benchmark.py

#!/usr/bin/env python3
"""
LNC Benchmark: Compare LNC=1 vs LNC=2 on Trainium2

Single standalone script that downloads a model, compiles it for both LNC
configurations, and benchmarks single-core + whole-instance throughput.

Usage:
    python3 lnc_benchmark.py              # Run full comparison
    python3 lnc_benchmark.py --lnc 2      # Run only LNC=2
    python3 lnc_benchmark.py --lnc 1      # Run only LNC=1

Requirements:
    pip install open_clip_torch

Instance: trn2.3xlarge (1 Trainium2 chip, 4 physical NeuronCores)
"""

import argparse
import json
import os
import subprocess
import sys
import time

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_NAME = "ViT-SO400M-14-SigLIP-384"
PRETRAINED = "webli"
IMAGE_SIZE = 384
BATCH_SIZE = 1
WARMUP_ITERS = 10
BENCH_ITERS = 50
MULTI_CORE_DURATION = 60  # seconds
COMPILER_ARGS = [
    "--optlevel", "3",
    "--model-type", "transformer",
    "--auto-cast", "matmult",
]

LNC_CONFIGS = {
    1: {"file": "siglip_lnc1.pt", "cores": 8},
    2: {"file": "siglip_lnc2.pt", "cores": 4},
}

# ---------------------------------------------------------------------------
# Worker script template (written to /tmp and run as subprocess)
# ---------------------------------------------------------------------------
WORKER_TEMPLATE = """
import torch, torch_neuronx, os, sys, time, json
os.environ["NEURON_RT_LOG_LEVEL"] = "ERROR"

core_id = int(sys.argv[1])
model_file = sys.argv[2]
duration = float(sys.argv[3])
warmup = int(sys.argv[4])
bench_iters = int(sys.argv[5])
batch_size = int(sys.argv[6])
image_size = int(sys.argv[7])

os.environ["NEURON_RT_VISIBLE_CORES"] = str(core_id)
model = torch.jit.load(model_file)
model.eval()
inp = torch.randn(batch_size, 3, image_size, image_size)

# warmup
for _ in range(warmup):
    with torch.no_grad():
        model(inp)

if duration > 0:
    # timed duration mode (multi-core)
    count = 0
    t0 = time.time()
    while time.time() - t0 < duration:
        with torch.no_grad():
            model(inp)
        count += 1
    elapsed = time.time() - t0
    print(json.dumps({"core": core_id, "throughput": count / elapsed}))
else:
    # fixed iteration mode (single-core)
    times = []
    for _ in range(bench_iters):
        t0 = time.time()
        with torch.no_grad():
            model(inp)
        times.append(time.time() - t0)
    avg = sum(times) / len(times)
    print(json.dumps({"core": core_id, "throughput": 1.0 / avg, "latency_ms": avg * 1000}))
"""

def _write_worker():
    """Write the worker script to /tmp."""
    path = "/tmp/_lnc_worker.py"
    with open(path, "w") as f:
        f.write(WORKER_TEMPLATE)
    return path

def _run_worker(worker_file, core_id, model_file, duration, lnc):
    """Launch a single worker subprocess."""
    env = os.environ.copy()
    env["NEURON_RT_VISIBLE_CORES"] = str(core_id)
    env["NEURON_RT_LOG_LEVEL"] = "ERROR"
    env["NEURON_LOGICAL_NC_CONFIG"] = str(lnc)

    args = [
        sys.executable, worker_file,
        str(core_id), model_file,
        str(duration), str(WARMUP_ITERS), str(BENCH_ITERS),
        str(BATCH_SIZE), str(IMAGE_SIZE),
    ]
    return subprocess.Popen(args, env=env, stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE, text=True)

def get_model():
    """Download and return the SigLIP vision encoder."""
    import open_clip
    print(f"Downloading {MODEL_NAME} ({PRETRAINED})...")
    model, _, _ = open_clip.create_model_and_transforms(
        MODEL_NAME, pretrained=PRETRAINED)
    model.eval()
    visual = model.visual
    visual.eval()
    print(f"  Vision encoder ready "
          f"({sum(p.numel() for p in visual.parameters()) / 1e6:.0f}M params)")
    return visual

def compile_model(lnc):
    """Compile the model for a given LNC setting. Skips if .pt exists."""
    import torch
    import torch_neuronx

    cfg = LNC_CONFIGS[lnc]
    out_file = cfg["file"]

    if os.path.exists(out_file):
        size_mb = os.path.getsize(out_file) / (1024 * 1024)
        print(f"  {out_file} already exists ({size_mb:.0f} MB) -- skipping")
        return

    visual = get_model()
    example = torch.randn(BATCH_SIZE, 3, IMAGE_SIZE, IMAGE_SIZE)
    args = COMPILER_ARGS + ["--lnc", str(lnc)]
    print(f"  Compiling LNC={lnc} (flags: {' '.join(args)})...")
    t0 = time.time()
    traced = torch_neuronx.trace(
        visual, example,
        compiler_workdir=f"./compile_lnc{lnc}",
        compiler_args=args,
    )
    elapsed = time.time() - t0
    torch.jit.save(traced, out_file)
    size_mb = os.path.getsize(out_file) / (1024 * 1024)
    print(f"  Saved {out_file} ({size_mb:.0f} MB, {elapsed:.0f}s)")

def bench_single_core(lnc, worker_file):
    """Benchmark on one NeuronCore via subprocess."""
    cfg = LNC_CONFIGS[lnc]
    p = _run_worker(worker_file, 0, cfg["file"], 0, lnc)
    stdout, _ = p.communicate(timeout=300)
    for line in stdout.strip().split("\n"):
        if line.startswith("{"):
            data = json.loads(line)
            return data["throughput"], data["latency_ms"]
    return 0.0, 0.0

def bench_multi_core(lnc, worker_file):
    """Benchmark all cores in parallel via subprocesses."""
    cfg = LNC_CONFIGS[lnc]
    procs = [_run_worker(worker_file, c, cfg["file"], MULTI_CORE_DURATION, lnc)
             for c in range(cfg["cores"])]

    per_core = {}
    for p in procs:
        stdout, _ = p.communicate(timeout=MULTI_CORE_DURATION + 120)
        for line in stdout.strip().split("\n"):
            if line.startswith("{"):
                data = json.loads(line)
                per_core[data["core"]] = data["throughput"]
    return sum(per_core.values()), per_core

def run_benchmarks(lnc):
    """Compile + benchmark for one LNC setting."""
    print(f"\n{'='*60}")
    print(f"  LNC={lnc}  ({LNC_CONFIGS[lnc]['cores']} logical cores)")
    print(f"{'='*60}")

    compile_model(lnc)
    worker_file = _write_worker()

    print(f"\n  Single-core benchmark "
          f"({WARMUP_ITERS} warmup, {BENCH_ITERS} timed)...")
    tp_single, lat = bench_single_core(lnc, worker_file)
    print(f"    Throughput: {tp_single:.2f} img/s")
    print(f"    Latency:    {lat:.2f} ms")

    n = LNC_CONFIGS[lnc]["cores"]
    print(f"\n  Whole-instance benchmark ({n} cores, {MULTI_CORE_DURATION}s)...")
    tp_total, per_core = bench_multi_core(lnc, worker_file)
    for c in sorted(per_core):
        print(f"    Core {c}: {per_core[c]:.2f} img/s")
    print(f"    Total: {tp_total:.2f} img/s")

    try:
        os.remove(worker_file)
    except OSError:
        pass

    model_size = os.path.getsize(LNC_CONFIGS[lnc]["file"]) / (1024 * 1024)
    return {
        "lnc": lnc, "cores": n,
        "single_throughput": tp_single, "single_latency_ms": lat,
        "total_throughput": tp_total, "per_core": per_core,
        "model_size_mb": model_size,
    }

def print_comparison(results):
    """Print a side-by-side comparison table."""
    print(f"\n{'='*60}")
    print("  COMPARISON: LNC=2 (default) vs LNC=1")
    print(f"{'='*60}")
    print(f"  Model: {MODEL_NAME}")
    print(f"  Instance: trn2.3xlarge\n")
    print(f"  {'Metric':<30} {'LNC=2':>12} {'LNC=1':>12}")
    print(f"  {'-'*54}")

    r2, r1 = results.get(2, {}), results.get(1, {})

    def fmt(key, label, unit=""):
        v2, v1 = r2.get(key), r1.get(key)
        s2 = f"{v2:.2f}{unit}" if v2 is not None else "---"
        s1 = f"{v1:.2f}{unit}" if v1 is not None else "---"
        print(f"  {label:<30} {s2:>12} {s1:>12}")

    def fmti(key, label):
        v2, v1 = r2.get(key), r1.get(key)
        s2 = f"{v2}" if v2 is not None else "---"
        s1 = f"{v1}" if v1 is not None else "---"
        print(f"  {label:<30} {s2:>12} {s1:>12}")

    fmti("cores", "Logical cores")
    fmt("model_size_mb", "Compiled model size", " MB")
    fmt("single_throughput", "Single-core throughput", " img/s")
    fmt("single_latency_ms", "Single-core latency", " ms")
    fmt("total_throughput", "Whole-instance throughput", " img/s")
    print(f"  {'-'*54}")

def _run_phase(lnc):
    """Run this script in a subprocess with the right LNC env var."""
    env = os.environ.copy()
    env["NEURON_LOGICAL_NC_CONFIG"] = str(lnc)
    env["NEURON_RT_LOG_LEVEL"] = "ERROR"
    cmd = [sys.executable, __file__, "--_run", str(lnc)]
    proc = subprocess.run(cmd, env=env, capture_output=True, text=True)
    sys.stdout.write(proc.stdout)
    if proc.stderr.strip():
        sys.stderr.write(proc.stderr)
    if proc.returncode != 0:
        sys.exit(proc.returncode)
    for line in reversed(proc.stdout.strip().split("\n")):
        if line.startswith("{"):
            return json.loads(line)
    return None

def main():
    parser = argparse.ArgumentParser(description="LNC Benchmark")
    parser.add_argument("--lnc", type=int, choices=[1, 2], default=None,
                        help="Run only this LNC config (default: both)")
    parser.add_argument("--_run", type=int, help=argparse.SUPPRESS)
    args = parser.parse_args()

    if args._run is not None:
        results = run_benchmarks(args._run)
        print(json.dumps(results))
        return

    print("=" * 60)
    print("  LNC Benchmark -- SigLIP on Trainium2")
    print("=" * 60)

    configs = [args.lnc] if args.lnc else [2, 1]
    all_results = {}
    for lnc in configs:
        r = _run_phase(lnc)
        if r:
            all_results[lnc] = r

    if len(all_results) == 2:
        print_comparison(all_results)

if __name__ == "__main__":
    main()