PyTorch Adaptation and Migration for Domestic Accelerators: From Operator Differences to Performance Alignment

Note: This is a review article (no first-hand experimental data), based on public adapter-framework documentation and migration practices.

Experiment Objective

Domestic accelerators (Ascend, Cambricon, MetaX, T-Head, etc.) are increasingly common in supercomputing and AI-compute centers, but their ecosystem maturity still lags CUDA. The goal: explain the full path for migrating a PyTorch program to a domestic accelerator — what the adapter layer actually does, which parts of your training script will "silently slow down" or "error out," and how to align performance with a migration checklist and baseline comparison. This is not "change one line"; it is a systematic retrofit grounded in understanding the adapter stack.

Code Walkthrough

Part 1: The three-layer adapter structure

Your training script (torch.nn / torch.optim / torch.distributed)
        │
        ▼
PyTorch core (original API unchanged)
        │
        ▼
Adapter layer: device plugin (torch_npu / torch_mlu / ...) + custom kernels
        │  │
        ├── operator registration: aten ops → accelerator kernels
        ├── device management: DeviceGuard, memory allocator, events/streams
        └── compile backend: graph mode (torch.compile / AI compiler) fusion
        │
        ▼
Accelerator driver + runtime (ACL / CNRT / vendor runtime)

Key points:

1. The training script barely changes: the adapter mounts into PyTorch as a "device plugin"; model.cuda() becomes model.to('npu') (or via a PYTORCH_DEVICE env var), and forward/backward/optimizer code stays the same — this is why migration cost is low.
2. Operator registration is the key difference: PyTorch's aten ops (matmul, conv, softmax, LayerNorm, …) have native kernels on CUDA; the adapter must supply its own kernel for each. Coverage decides compatibility — uncovered ops fall back to CPU, showing as "runs but extremely slow."
3. Graph mode is the performance catch-up tool: early domestic kernels lag CUDA libraries in quality, so torch.compile/AI compilers fuse many ops into one large kernel, cutting launch overhead and memory round-trips — often recovering 20%-40% of the gap.

Part 2: Step one — device-agnostic refactor

# Before: hardcoded cuda
device = 'cuda'
model = model.to(device)

# After: device-agnostic
import os
device = os.environ.get('AI_DEVICE', 'cuda')   # supports cuda/npu/mlu/cambricon
model = model.to(device)

# Key compatibility points:
# 1. tensor memory API: tensor.cuda() → tensor.to(device)
# 2. distributed backend: dist.init_process_group(backend='nccl') → adapter's (hccl/cncl)
# 3. mixed precision: some torch.cuda.amp APIs need adapter equivalents
# 4. custom CUDA kernels (.cu source) cannot migrate directly → rewrite or use op library

Explanations:

• to(device) is the universal entry: all memory allocation and device switching go through it; hardcoded cuda is the first place to audit.
• Distributed backends differ: NCCL becomes the vendor collective library (Ascend HCCL, Cambricon CNCL); init_process_group parameters change, but PyTorch's DDP upper logic stays.
• Mixed-precision APIs differ: torch.cuda.amp.GradScaler has equivalents in some adapters, but the AMP operator whitelist differs — verify against the adapter's docs.
• Custom CUDA kernels are the largest migration cost: any hand-written .cu operator must be rewritten with the adapter's operator library or with Triton (chapter 08's Triton article pays off here).

Principles

Why "it runs after changing device" ≠ "migration complete"

Three levels of pitfalls, shallow to deep:

1. Compatibility layer (it runs): APIs callable, operators covered. But CPU-fallback operators can slow a whole section 10× — runs ≠ trains.
2. Performance layer (it trains): operator coverage plus graph fusion, training speed approaches the CUDA baseline. Acceptance: end-to-end throughput gap < 20%.
3. Precision layer (it compares): floating-point implementation differences (different FlashAttention implementations, accumulation order) can shift the loss curve slightly — verify with same seed and loss alignment.

Performance alignment method

Baseline: same model, same data on CUDA; record throughput/memory/per-step time
Migration: same config on the accelerator
Gap localization: per-op profiling (adapter tools / torch.profiler)
Optimization: op fusion (torch.compile) → allocator tuning → collective tuning → data pipeline
Acceptance: throughput gap ≤20%, loss curves aligned, 48h stability

Key insight: the gap is usually not "one slow operator" but the combination of "several fallback ops + no graph fusion + different memory allocation." Fix fallback first (worst), then graph mode, then memory/communication params — this order gives the most return.

Results

(Estimated, not measured; varies widely by vendor/model)

Migration stage	Throughput vs CUDA baseline	Main issue
device change only	20%-50%	op fallback, no fusion
+ fix fallback ops	60%-80%	per-op kernel efficiency gap
+ torch.compile fusion	80%-100%	near baseline; some ops slow
+ memory/comm tuning	90%-110%	on par or slightly better

Engineering Takeaways

1. Audit operator coverage before migration: scan the op list your script uses with adapter tools and compare against the support table; list fallback ops as the schedule's main work item — this is the primary basis for estimation.
2. Custom CUDA kernels are the watershed: projects with hand-written ops double the migration cost; pure-PyTorch projects usually reach "trains" in 1-2 weeks.
3. Acceptance criteria belong in contracts/deliverables: throughput gap, loss alignment, and stability are non-negotiable — "it runs" is not a delivery standard.
4. Pair with heterogeneous scheduling: chapter 07's heterogeneous-resource-management covers "how to manage mixed deployment"; this article covers "how to migrate one job" — in a mixed cluster, domestic cards cannot run CUDA ops, so migration quality directly decides cluster utilization.

Next article preview: single-job migration done; back to global training stability — "Pretraining Stability Engineering: Loss-Spike Detection and Checkpoint Restart Strategies."