Conversation
github-actions
Bot
added
the
community-contribution
for more information, see https://pre-commit.ci
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| bool ReleaseWorker::process_remap_batch(const std::shared_ptr<RemapHookContext> &context, | ||
| bool requeued) { | ||
| RemapWorker::process(context, /*serial_requeue=*/requeued); | ||
| return context->done.load(std::memory_order_acquire) != 0; |
There was a problem hiding this comment.
In serial-driver mode, the first remap attempt uses requeued=false. If its release callback has not yet queued the matching release, the remap waits for that release instead of returning. Because the same worker must process the release, it can never make progress, permanently blocking this remap and subsequent waits. Every serial-worker remap attempt must remain requeueable while its release dependency is pending.
| void close() { | ||
| drain_remap(); | ||
| drain_pending(); | ||
| if (!reserved_) return; | ||
| if (mapped_) { | ||
| check_cuda(cuMemUnmap(address_, bytes_), "cuMemUnmap(close)"); | ||
| mapped_ = false; | ||
| check_cuda(cuMemRelease(handle_), "cuMemRelease(close)"); | ||
| handle_ = 0; | ||
| } | ||
| check_cuda(cuMemAddressFree(address_, bytes_), "cuMemAddressFree"); | ||
| reserved_ = false; | ||
| address_ = 0; | ||
| } |
There was a problem hiding this comment.
close() waits only until an asynchronous H2D has been submitted, not until the GPU copy has completed, and then immediately unmaps and releases its destination. If a caller closes an allocation after remap_and_copy_after returns but before the copy stream completes, the GPU can continue writing to unmapped memory, causing a device fault or a corrupted transfer. The close and destructor paths must retain the mapping until the relevant completion event finishes.
| if self.device.type != "cuda": | ||
| raise ValueError(f"CUDA VMM allocation requires a CUDA device, got {self.device}") |
There was a problem hiding this comment.
The advertised MUSA path is unreachable through this Python API: construction rejects every device type except cuda, and _raw_cuda_stream checks only cuda_stream despite the stated musa_stream fallback. On a MUSA runtime without a complete torch.cuda alias, callers cannot construct or schedule this feature even if the native source is translated successfully.
| std::shared_ptr<ReleaseHookContext> release_hooks_after( | ||
| std::vector<std::shared_ptr<VMMActivationSlot>> slots, uintptr_t raw_stream) { | ||
| RECORD_USER_SCOPE("vmm::release_hooks_after"); | ||
| // Ensure the worker thread exists before any host func can fire. | ||
| ReleaseWorker::instance(); | ||
| auto context = std::make_shared<ReleaseHookContext>(); | ||
| context->tls_state = std::make_shared<at::ThreadLocalState>(); | ||
| for (auto &slot : slots) { | ||
| if (slot->pending_remap_) { | ||
| throw std::runtime_error("VMM batch release hook got a slot with an in-flight remap"); | ||
| } | ||
| if (slot->pending_) { | ||
| throw std::runtime_error("VMM batch release hook got a slot with an in-flight release"); | ||
| } | ||
| context->requests.push_back(slot->make_release_request()); | ||
| if (context->requests.back().handle == 0) { |
There was a problem hiding this comment.
The batch release helper accepts the same slot more than once, creating duplicate requests for one virtual address and allocation handle. The worker then attempts cuMemUnmap and cuMemRelease twice, records an asynchronous error on the shared context, and leaves the slot bookkeeping stale after the first release succeeds. Reject duplicate slots before installing the release context, as the batch remap helpers already do.
| void remap_slot_launch_h2d(const std::shared_ptr<RemapHookContext> &context, size_t slot_index, | ||
| uintptr_t raw_stream) { | ||
| RECORD_USER_SCOPE("vmm::remap_slot_launch_h2d"); | ||
| if (!context->defer_copy) { | ||
| throw std::runtime_error("VMM slot H2D launch requires a deferred-copy context"); | ||
| } | ||
| if (slot_index >= context->requests.size()) { |
There was a problem hiding this comment.
A deferred-copy context allows launch_slot_h2d to be called repeatedly for the same request because copy_submitted is never checked. Each call queues another copy and event record, but lifetime retention is installed only when h2d_launched first equals the request count. After adoption and reference release, the first callback can therefore destroy the event and pinned host storage before a later duplicate copy completes. Reject a launch after that request has already been submitted.
2 participants
Description
Add a fixed-VA CUDA VMM activation slot so a Graph-captured tensor address can keep its virtual address while physical backing is released and remapped outside Graph replay. This is the TE primitive behind local whole-layer CUDA Graph + fine-grained CPU activation offload.
Fixes # (issue)
Local CUDA Graph and CPU offload currently conflict on pointer semantics:
A naive
empty()reload allocates a new pointer. Capture recordedsrc.data_ptr(), so replay either reads stale storage or silently aliases an allocator reuse of that VA. Related: NVIDIA/Megatron-LM#3697.This PR does not implement Megatron scheduling. It exposes the VMM slot, resident workers, and deferred H2D so a caller (Megatron
fine_grained_activation_offload) can:slot.tensor.data_ptr()(stable for the run).CUDA driver names stay CUDA-shaped in source. On MUSA they map at compile time through
musify.h(cuMemMap→muMemMap,cudaMemcpyAsync→musaMemcpyAsync, …). Python keeps CUDA names (CUDAActivationVMMAllocation,cuda_stream); a MUSA runtime still works whentorch.cudais aliased, anddevice.type in ("cuda", "musa")covers TE tests that import this module withoutmusa_patch.Type of change
Changes
transformer_engine/pytorch/csrc/extensions/vmm_activation.cppand bind it frompybind.cpp/extensions.h.transformer_engine/pytorch/vmm_activation.py.tests/pytorch/test_vmm_activation.py(async reload / host-callback lifetime).musify.hVMM aliases used by this file (cuMemAddressReserve,cuMemCreate,cuMemMap,cuMemSetAccess,cuMemUnmap,cuMemRelease,cuMemGetInfo, …).Not in this PR’s intended review surface: GEMM / DPA /
cpu_offload.pyMegatron handler changes that happen to sit on the same working branch.Why a TE primitive (not a Megatron allocator wrapper)
Graph capture records a device pointer. CPU offload must drop physical pages. Those two facts only compose if the Graph sees a VA that outlives the pages mapped under it:
The captured Graph never sees a new pointer. Physical handles may change; the VA must not.
CUDAActivationVMMAllocationconstructs the slot, wraps a non-owningfrom_blobview, and assertstensor.data_ptr() == info["address"]after construct / remap / adopt. Using that view while the slot is unmapped is a hard page fault — that is intentional; it turns the silent allocator-reuse bug into a crash.Slot lifecycle
Synchronous
unmap_and_release()/create_and_remap()remain for tests and fallback. Training uses the async path.Address stability check after every remap / adopt:
Python API
transformer_engine.pytorch.vmm_activation:CUDAActivationVMMAllocation(shape, stride, dtype, device)release_hooks_after(allocs, stream)remap_hooks_after(allocs, stream)remap_and_copy_after/remap_and_copy_slot_afterremap_only_slot_afterlaunch_remap_slot_h2d(ctx, i, stream)remapped, submit H2D onstreamwait_remap_slot_on_stream(ctx, i, stream)event_recorded,streamwaits on slot eventwait_until_remap_slot_submittedwait_remap_copy_on_stream/enqueue_remap_copy_waitvmm_driver_memory_info()cuMemGetInfofree/total (caching-allocator stats do not covercuMemCreatepages)vmm_enable_trace/vmm_initialize_workers/vmm_set_serial_driver_workersStream handles:
stream.cuda_stream, withmusa_streamfallback.Safety checks on remap submit: unique slots, reserved VA, no in-flight remap, mapped slot must already have a pending release, host tensor pinned / offset 0 / large enough. Adopted mappings are never unmapped by context cleanup.
Debug:
MEGATRON_VMM_REMAP_DEBUG=1prints per-call remap/release stages (cuMemUnmap,cuMemSetAccess, remap wait, H2D launch wait).Files
Checklist:
Suggested coverage:
data_ptr()unchanged across unmap/remap and deferred H2D.release_hooks_after+remap_and_copy_afterrepeated reload (test_repeated_async_reload_releases_context_off_host_callback): no host-callback deadlock, values match, contexterror==0.remap_only_slot_afterdoes not submit H2D;launch_remap_slot_h2ddoes; compute stream waits on the slot event.cuMemGetInfo(not PyTorchmemory_allocated) when claiming physical pages returned.