GB10/Blackwell: training hangs at sample-tensor pre-upload (sample 0/N stalled, no progress)

[dndungu]

Symptom

On NVIDIA GB10 Blackwell (aarch64, unified memory), the GPU training path hangs silently and permanently during the bulk sample-tensor + graph-parameter pre-upload step. The operator observes:

• Pod remains phase=running in Spark indefinitely — no log output after the pre-upload init line (see Stall location below).
• nvidia-smi hangs on the host.
• DELETE /api/v1/pods/{name} via the Spark HTTP API times out (>180 s) — host-side podman is wedged.
• No error is returned; no panic; no CUDA error code surfaced.
• The issue is reproducible across 6 independent attempts spanning ztensor v1.6.0 and earlier revisions (see What we've tried).

Hardware / Software

Field	Value
GPU	NVIDIA GB10 Blackwell
Architecture	aarch64 (linux/arm64)
Memory model	Unified memory (no discrete VRAM)
ztensor	v1.6.0
Zerfoo	github.com/zerfoo/zerfoo v1.48.1-0.20260428224953-7539627aad66 (HEAD at time of hang)
Wolf commit	35087d91 (training image localhost:5000/wolf-train-crossasset:t98-35087d91)
CUDA graph capture	Enabled (default in ztensor v1.6.0; ztensor#93 resolved the prior capture hang)

Reproduction

Submit the following Spark Pod manifest to http://<dgx-host>:8080/api/v1/pods:

apiVersion: v1
kind: Pod
metadata:
  name: wolf-train-crossasset-gpu-35087d91
spec:
  priorityClassName: high
  restartPolicy: Never
  containers:
    - name: train
      image: localhost:5000/wolf-train-crossasset:t98-35087d91
      args:
        - -gpu
        - -bars=/data/experiments/coin-all-1m-bars.csv
        - -epochs=50
        - -folds=5
        - -output=/models/crossasset-gpu-results.json
        - -save=/models/coin-crossasset-gpu.zcam
      env:
        - name: LD_LIBRARY_PATH
          value: /usr/local/cuda/lib64:/opt/zerfoo/lib
      resources:
        limits:
          nvidia.com/gpu: "1"
          cpu: "6"
          memory: "24Gi"
        requests:
          cpu: "6"
          memory: "16Gi"

Single-shot submit command:

curl -X POST http://192.168.86.250:8080/api/v1/pods \
  --data-binary @deploy/spark/train-crossasset-gpu.yaml \
  -H 'Content-Type: application/yaml'

The caller's trainWithResult logic (Wolf internal/crossasset/crossasset.go) calls uploader.UploadWeights(f32) where uploader is the ztensor GPU engine, passing 106,652 sample tensors plus ~50 autograd graph parameter tensors. The call never returns.

Stall location

Last log line printed before the hang (timestamp 23:36:13, no further output after 36+ minutes):

trainWithResult: T4.2 pre-uploading 106652 sample tensors + 50 graph params to GPU

This is the log line immediately before uploader.UploadWeights(f32) is called. The call enters ztensor's GPUEngine.UploadWeights (ztensor/compute/gpu_engine.go) and never returns.

What we've tried (6 attempts, same symptom family)

Attempt	Change from prior	Outcome
#1–#5	Various: CUDA graph disabled (ZERFOO_DISABLE_CUDA_GRAPH=1), different epoch/fold counts, different memory limits	Silent stall past "Using GPU engine" — hang before trainWithResult even started
#6 (2026-04-28)	ztensor v1.6.0 (reverted from v1.7.0 via PR #141), ztensor#93 CUDA-graph-capture fix in place, T9.12 softmax fix applied	Stall advanced one stage deeper — now inside trainWithResult at the UploadWeights call — but same fingerprint: 0 GPU progress, no error, host wedged

Prior to attempt #6, attempts #1–#5 stalled earlier (before trainWithResult) due to a separate CUDA graph capture deadlock resolved by ztensor#93. That fix was necessary but not sufficient.

Workaround in use

None on the GPU path. Plan B is CPUEngine[float64] (Wolf T9.22), which bypasses UploadWeights entirely.

What we'd need to investigate further

1. nsys trace — Per Wolf lore entry L-0008, nsys --trace=cuda itself deadlocks GB10 GPU training (CUPTI listener conflict with CUDA graph capture). If ztensor v1.6.0 CUDA graph capture is active during UploadWeights, nsys profiling may not be viable without first disabling capture for the upload phase.
2. CUDA driver / runtime versions on GB10 — exact libcuda.so version and CUDA runtime version on the aarch64 host would help narrow whether this is a driver-level unified-memory page-migration stall or a CUDA API hang.
3. GB10/Blackwell-specific UploadWeights constraints — any known restrictions on: maximum single-allocation size, tensor alignment requirements for unified-memory H2D bulk copies, or interaction between CUDA graph capture and large cudaMemcpy calls on GB10.
4. Does disabling CUDA graph capture for the upload phase unblock it? — If UploadWeights is called while the CUDA graph capture context is active, the memcpy may be recorded rather than executed, causing the hang. A flag or API to temporarily exit capture mode during bulk upload would be a minimal fix to test.

Cross-references

• Wolf PR #176: https://github.com/feza-ai/wolf/pull/176 — doc-only stall report for T9.8 attempt chore(main): release 0.4.0 #6
• Wolf lore L-0008: nsys --trace=cuda deadlocks Wolf GPU training on GB10 (CUPTI/CUDA-graph conflict)
• Wolf ADR-062: PyTorch speed parity strategy (context for the GPU training goal)
• ztensor#93: prior CUDA graph capture hang on GB10 (resolved, but insufficient for this stall)
• ztensor PR #141: revert of ztensor v1.7.0 (separate deadlock)

[dndungu]

Root cause: 106,652 individual cudaMalloc + cudaMemcpy calls in a tight loop

compute/gpu_engine.go::UploadWeights (lines 350-582) iterates over the input slice and, for each non-quantized float32 tensor, executes:

devPtr, err := e.allocWeight(byteSize)        // -> e.runtime.Malloc(byteSize)
src := unsafe.Slice((*byte)(unsafe.Pointer(&data[0])), byteSize)
if err := e.uploadBytes(devPtr, src); err != nil { ... }  // -> e.runtime.Memcpy H2D

With managedMem == false (the default; ZERFOO_ENABLE_MANAGED_MEM is unset in the reporter's manifest) and CUDA-graph capture inactive at upload time, this is one synchronous cudaMalloc plus one synchronous cudaMemcpy(HostToDevice) per input tensor.

Wolf's caller passes 106,652 sample tensors + 50 graph parameter tensors as a single slice. That is ~213,000 round-trips through the CUDA driver back-to-back, with no batching, no stream sharding, and no host-side coalescing.

Why this wedges GB10 specifically

On GB10 Blackwell with unified memory:

1. Every cudaMalloc creates a managed allocation tracked in the kernel page table.
2. The driver serializes page-table updates against the GPU command queue.
3. Each cudaMemcpy H2D in non-capture mode synchronously waits for the prior op on the default stream.
4. After ~few thousand iterations the driver's per-context allocation table contention plus the default-stream queue pressure stalls the host runtime entirely. nvidia-smi hangs because the user-mode driver waits on the same lock the upload loop is hammering. podman can't reap because the container's CUDA context is wedged.

This matches the reported fingerprint exactly:

• "Sample 0/N stalled, no progress" -- the loop is making forward progress (mallocs are completing) but the host has no observability because logging is below the per-iteration granularity.
• nvidia-smi hangs -- driver-level lock contention.
• DELETE times out -- podman cannot tear down a wedged CUDA context until cudaMalloc returns.
• No error -- nothing failed; everything is blocked waiting on serialized driver work.

Why ztensor#93 was insufficient

ztensor#93 fixed cudaMallocManaged calls inside a CUDA-graph-capture region (which are illegal). That makes the capture-aware path safe. But the reporter's path is the non-capture path -- capture has not begun yet at trainWithResult upload time -- so the #93 routing doesn't apply. The bug is in the per-iteration synchronous Malloc+Memcpy pattern, not in graph capture.

Workaround (caller side, no ztensor change required)

Concatenate sample tensors into one or a few large tensors before calling UploadWeights:

// Wolf: replace per-sample uploads with a single bulk upload.
total := 0
for _, t := range samples {
  total += len(t.Data())
}
bulk := make([]float32, total)
offsets := make([]int, 0, len(samples))
off := 0
for _, t := range samples {
  copy(bulk[off:], t.Data())
  offsets = append(offsets, off)
  off += len(t.Data())
}
bulkTensor := tensor.New(...) // wrap `bulk` in a single TensorNumeric[float32]
uploader.UploadWeights([]*tensor.TensorNumeric[float32]{bulkTensor})
// then: keep `offsets` and slice into the GPU buffer per (epoch, batch, sample).

This collapses 106,652 driver round-trips into 1. Should resolve the stall without any ztensor change.

Fix (ztensor side)

Add a bulk path in UploadWeights for the F32 / FP16 / BF16 branches when the input slice exceeds a threshold (e.g., 1024 tensors):

1. First pass: compute total byte size across all eligible tensors.
2. One e.allocWeight(total) for the combined region.
3. One e.uploadBytes(devPtr, combinedSrc) (or, on managed-memory hosts, one CPU-side copy() into the managed slice).
4. Second pass: issue tensor.NewGPUStorageFromPtr for each input tensor as a sub-slice view of the single allocation.

This preserves the existing per-tensor API while collapsing the driver round-trips. Effectively the same pattern preAllocateWorkspaces (line 625) already uses elsewhere in the engine.

Verification path

Once a fix is in:

1. Build a test that uploads 100K small float32 tensors and asserts wall-clock < (say) 5s.
2. Run on GB10 via Spark with the reporter's t98-35087d91 config -- training should advance past the upload line.

Will leave this issue open; the workaround unblocks Wolf in the meantime, but the underlying ztensor pattern is hostile to large training batch counts and should be fixed in the engine.

[dndungu]

Verified live on GB10

v1.8.0 (PR #104) plus Wolf bump (feza-ai/wolf #178, image localhost:5000/wolf-train-crossasset:44fb4790) was run end-to-end on the DGX Spark host:

Submit:    POST /api/v1/pods (ztensor103-verify-r3)
Args:      -gpu -bars=/data/experiments/coin-all-1m-bars.csv -epochs=2 -folds=2
Image:     localhost:5000/wolf-train-crossasset:44fb4790
Pod state: STATUS=running 06:45:23 -> STATUS=completed 07:03:55 (18m32s wall-clock, exit 0)

Logged progress (no hang, no driver wedge):
  Fold 0: acc=0.6436 (train=16408, test=8204, time=424.5s)
  Fold 1: skipped (only 1 test samples)        # Wolf-side -folds=2 split artifact, not ztensor
  Mean accuracy: 0.6436 +/- 0.0000

UploadWeights returned cleanly. Training advanced through both epochs of fold 0 and produced real accuracy. The previous symptom (no progress, nvidia-smi hangs, DELETE times out) is gone.

Closing.