HCIP-AI EI Developer: One PyTorch Codebase for GPU Workstations and Ascend Production

One PyTorch Codebase for GPU Workstations and Ascend Production

A computer-vision team maintains a single PyTorch training repository used on developer GPU workstations (RTX 3090) and on a production ModelArts cluster with Ascend 910. The code hard-codes model.cuda(), torch.cuda.amp.autocast(), and DataLoader(num_workers=16, pin_memory=True). On Ascend the job either raises a no-CUDA-device error or silently runs on CPU at about 40 images/s. A second problem: a custom __getitem__ that opens JPEGs with PIL and applies Python-level augmentation consumes about 65% of each training step on both platforms. The team proposes adding if is_npu: branches throughout the code and installing a separate fork for production. Constraints: one maintained branch, no model architecture change, and a 20% step-time reduction target. The team has profiled only the GPU path so far.

1. Evaluate the proposed `if is_npu` fork. What is the better approach, and why?

Read the worked answer

A fork plus scattered device branches solves the immediate crash but creates two code paths that will drift. Every bug fix, augmentation change, and checkpoint format must be replicated, and the production path will be tested less often than the GPU path. That is a poor trade for a device-placement problem.

A thin device-abstraction layer is better:

Import the Ascend extension for PyTorch before device use, then select device = torch.device('npu:0') when running on Ascend and cuda:0 otherwise.
Replace .cuda() with .to(device) everywhere, including buffers and loss tensors.
Wrap AMP behind one helper: CUDA autocast on GPU, the torch_npu AMP API with loss scaling on Ascend.
Set pin_memory=False on Ascend, where the CUDA pin-memory path does not apply.
Keep if statements at the boundary layer only, not inside the model or training loop.

The judgement is that device portability is a configuration concern, not a model concern. A single branch with one compatibility module is auditable; a fork is not.

2. Where would you focus to hit the 20% step-time target? Explain how profiling changes the decision.

Read the worked answer

The stated 65% share for Python preprocessing is the strongest signal, but it is a GPU-path measurement and must be confirmed on Ascend before committing effort.

If profiling confirms the data pipeline dominates:

Replace per-image PIL decode plus Python augmentation with a compiled pipeline: torchvision v2 transforms with batched ops, a DALI-style loader, or pre-decoded binary shards.
Do not simply raise num_workers; if the work is GIL-bound in Python, more workers add memory and context switching without proportional gain.
Overlap host preprocessing with device compute using prefetch and non_blocking transfers.

If profiling instead shows the NPU step is compute-bound, the 20% target cannot come from the loader. Then the options are mixed precision, operator fusion, or a more efficient input resolution, not more workers.

The decision rule: optimize the measured dominant term on the target platform, then re-profile. A GPU-only profile can rank bottlenecks differently because the NPU step time, host-device copy behavior, and AMP implementation all differ. Accept the change only if end-to-end images/s improves by 20% on Ascend without changing convergence.

3. What checks must pass before declaring the unified codebase production-ready on Ascend?

Read the worked answer

I would require evidence in five areas:

Device audit: no tensor or module remains on CPU inadvertently. A short run that asserts every parameter and input batch is on the NPU catches silent fallbacks.
Numerical parity: with a fixed seed and the same data order, loss after 100 steps should match the GPU run within a stated tolerance. Mixed precision should be tested with loss scaling on and off.
AMP behavior: overflow detection and scale updates are active; a forced overflow test shows the scaler skips the step rather than producing NaNs.
Data determinism: the same sampler seed produces the same batches on both platforms; shuffling and worker seeding are not device-dependent.
Checkpoint portability: a checkpoint saved on Ascend loads on GPU and vice versa, and resume restores optimizer and scaler state.

A short benchmark on the production flavor is also needed, because a workstation pass does not prove distributed HCCL initialization or OBS data access. Passing these checks is what justifies retiring the fork proposal.