[AI][Bug][Ascend] Peak reset clears live allocator counters

[yswang777]

Issue Type

• Bug Report
• Feature Request
• Operator Implementation
• Platform Support
• Performance Issue
• Documentation

AI Agent Information

• Agent: OpenAI Codex
• Model: GPT-5-based Codex
• Session Context: Investigating incorrect Torch-FL peak allocated memory reported by the Qwen-Image-2.1 Ascend benchmark.

Summary

torch.flagos.reset_peak_memory_stats() erases current allocated and reserved bytes, as well as cumulative allocation counters, on a non-delegating allocator. A live 4 MiB tensor is reported as zero bytes immediately after resetting peaks. This makes subsequent peak readings omit live model weights and can corrupt allocation accounting when the live block is freed.

Environment

• Platform: Ascend 910C, CANN 9.0.0, driver 25.5.0
• Python: 3.11.15
• PyTorch: 2.10.0+cpu with Torch-FL PrivateUse1 backend
• Torch-FL: 0.1.0 installed Ascend build; the same reset implementation remains on main at e31df225678184c1a2f24370607fcc6ca7715dce
• Build configuration: FLAGOS_ACCELERATOR=ascend FLAGOS_BUILD_FLAGGEMS=1 FLAGOS_BUILD_FLAGGEMS_CPP=0 FLAGOS_BUILD_VENDOR=1
• Runtime configuration: one visible Ascend chip, FLAGOS_USE_CACHING_ALLOCATOR=1 (default)

Reproduction

import torch
import torch_fl

torch_fl.flagos.init()
live = torch.empty(1024 * 1024, dtype=torch.float32, device="flagos:0")
before = torch_fl.flagos.memory_stats(0)
torch_fl.flagos.reset_peak_memory_stats(0)
after = torch_fl.flagos.memory_stats(0)
print("before", before)
print("after", after)
assert after["allocated_bytes"] == before["allocated_bytes"]
assert after["peak_allocated_bytes"] == after["allocated_bytes"]

Expected vs Actual Behavior

Expected: allocated_bytes remains 4,194,304; reserved_bytes remains 20,971,520; both peak values reset to the corresponding current values. The allocation counters remain unchanged.

Actual output from the pre-fix Ascend build:

before {'allocated_bytes': 4194304, 'reserved_bytes': 20971520, 'peak_allocated_bytes': 4194304, 'peak_reserved_bytes': 20971520, 'num_alloc_calls': 1, 'num_free_calls': 0, 'num_device_malloc': 1, 'num_device_free': 0, 'num_alloc_retries': 0}
after {'allocated_bytes': 0, 'reserved_bytes': 0, 'peak_allocated_bytes': 0, 'peak_reserved_bytes': 0, 'num_alloc_calls': 0, 'num_free_calls': 0, 'num_device_malloc': 0, 'num_device_free': 0, 'num_alloc_retries': 0}

The assertion fails because 0 != 4194304. This is a wrong-value result, not an exception thrown by the runtime.

Root Cause Analysis

On backends that do not delegate to a vendor caching allocator, CachingDeviceAllocator::reset_stats assigns state.stats = AllocatorStats{}. The public Python method is named reset_peak_memory_stats, but this assignment clears current usage and cumulative counts too. The delegated branch calls the vendor's peak reset and is not implicated.

Proposed Solution

Under the existing device-state lock, set peak_allocated = bytes_allocated and peak_reserved = bytes_reserved. Do not modify live usage or cumulative counters. Leave the delegated caching-allocator branch unchanged.

Verification Plan

• Minimal old-build reproduction above.
• Integration regression with a live allocation, repeated reset, a second allocation and free, reserved bytes, and unchanged counters.
• Full allocator integration suite on Ascend.
• Unit tests, lint, and review of other allocator backends.

Context & Investigation

Read the public Python memory API, Python-to-C++ binding, allocator state updates and both branches of reset_stats. Ran the 4 MiB live-tensor reproducer on a single Ascend chip. Inspected current main and checked issue/PR search results for reset_peak_memory_stats and allocator peak reset; no matching report was found. Related DCU memory API issue #331 concerns shim availability rather than this allocator-counter reset.

Related Code Locations

• csrc/runtime/allocator/caching_device_allocator.cc:505: the erroneous whole-structure reset.
• torch_fl/csrc/module.cc:461: public reset binding.

Checklist - AI Agents MUST Complete All

• I have provided complete environment information
• I have included a minimal, self-contained reproducer
• I have included full wrong-value output (no exception is thrown)
• I have analyzed the root cause (not just symptoms)
• I have proposed a specific solution with implementation approach
• I have identified affected code locations with line numbers
• I have described how to verify the fix
• I have checked for duplicate issues
• All text is in English
• Code follows project conventions (checked existing code style)