fix(ascend): harden disk delta checkpoint handling

[wangx700]

What

Hardens disk-delta checkpoint handling for the Ascend disk weight-sync path, on top of the NPU platformization work in #413:

vime/utils/disk_delta.py

• Deterministic shard discovery (sorted *.safetensors) and fail-fast FileNotFoundError on empty checkpoint dirs
• Strict safetensors header validation: truncation checks on the length prefix and header bytes, header length bounded by file size, JSON-object header check
• Per-tensor data_offsets bounds checking (0 <= begin <= end <= data_size) and truncated-read detection in the tensor reader - corrupted or partially written checkpoints now raise a descriptive RuntimeError instead of silently materializing wrong weights

vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py

• Pinned host-buffer pool extracted into _make_pinned_pool, bounded to 8 GiB total pinned allocation, with graceful fallback to pageable .cpu() copies
• Buffer ownership hardened: pinned buffers are always returned to the pool via try/finally, both when a pool worker fails and when task submission is interrupted (no pinned-memory leak / pool starvation on transient errors)

Why

Mirrors the checkpoint hardening from #422 (bound disk delta buffers and validate checkpoints) onto the #413 platformization base, where the disk delta path runs on shared filesystems that can surface truncated or partially committed checkpoint files.

Notes

• Stacked on [NPU] Sync Ascend with main and consolidate NPU adaptations #413: the diff vs ascend includes [NPU] Sync Ascend with main and consolidate NPU adaptations #413's commits until that PR merges.
• Scope is intentionally minimal - only the two hardening-related files change.

Test plan

• pytest tests/unit (CPU)
• 910 end-to-end smoke: --update-weight-mode delta --update-weight-transport disk (in progress on the NPU environment)

[mergify]

⚠️ The sha of the head commit of this PR conflicts with #428. Mergify cannot evaluate rules on this PR. Once #428 is merged or closed, Mergify will resume processing this PR. ⚠️

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34549759 | 📁 Comparing 8f3ffc9 against latest (ce92eff)

  🔍 Preview build  

48 files changed · + 3 added · ± 45 modified

+ Added

• get_started/NPU.html
• _examples_synced/retool/README.html
• _examples_synced/search-r1/README.html

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 35 more...

[gemini-code-assist]

Code Review

This pull request introduces significant updates to the Vime framework, including migrating the CI pipeline to Buildkite, adding support for AMD ROCm 7.0.2, updating the CUDA Dockerfile with newer dependencies (FlashAttention, TransformerEngine, DeepGEMM, DeepEP), and implementing advanced features like Delta Weight Sync, On-Policy Distillation (OPD), and MemAgent long-context RL. The review feedback highlights several critical areas for improvement: resolving a potential NaN propagation bug in custom_convert.py when n_samples_per_prompt is 1, making the patch application in swe.py more robust by resetting the workspace between fallbacks, removing an unnecessary 10-second sleep bottleneck in generate.py, adding retries to HTTP requests in eval_ruler_hqa.py, and robustly handling 'tool' messages in token_delta.py. Additionally, the Dockerfiles should be hardened by verifying sed replacements, parameterizing MAX_JOBS to prevent OOMs, and aligning dependency versions between CUDA and ROCm.

[gemini-code-assist]

When n_samples_per_prompt is 1 (or if the number of samples falls back to 1 in the else branch), rewards_tensor will have a dimension of size 1 along dim=-1. In PyTorch, calling std() on a dimension of size 1 with the default unbiased=True returns NaN. This NaN will propagate through the division, corrupting the rewards and leading to NaN gradients during training. To prevent this, either use unbiased=False or guard the division by checking if the dimension size is greater than 1.

Suggested change

	if getattr(args, "advantage_estimator", None) in ["grpo", "gspo"] and getattr(
	args, "grpo_std_normalization", False
	):
	std = rewards_tensor.std(dim=-1, keepdim=True)
	rewards_tensor = rewards_tensor / (std + 1e-6)
	if getattr(args,

[gemini-code-assist]

Collapsing the patch application commands into a single shell execution with || is extremely fragile. If the first command (git apply --3way) partially applies some changes but ultimately fails, it leaves the working directory in a dirty/partially-modified state. The subsequent fallback commands (git apply or patch) will then run on top of this dirty state, leading to unexpected conflicts, patch corruption, or false-negative evaluation failures. To ensure correctness, the working directory should be reset to a clean state before each fallback attempt.

Suggested change

	# First-success-wins ladder collapsed into one exec (one sandbox round-trip).
	ladder = " || ".join(
	f"({cmd})"
	for cmd in (
	f"git apply --3way --whitespace=nowarn {_PATCH}",
	f"git apply --whitespace=nowarn {_PATCH}",
	f"patch -p1 --no-backup-if-mismatch < {_PATCH}",
	)
	)
	ec, _, _ = await ev.exec(f"cd {workdir} && ({ladder})", user="agent", check=False, timeout=120)
	return ec == 0
	for cmd in (
	f"git apply --3way --whitespace=nowarn {_PATCH}",
	f"git apply --whitespace=nowarn {_PATCH}",
	f"patch -p1 --no-backup-if-mismatch < {_PATCH}",
	):
	await ev.exec(f"cd {workdir} && git reset --hard && git clean -fd", user="agent", check=False, timeout=60)
	ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="agent", check=False, timeout=120)
	if ec == 0:
	return True
	return False

[gemini-code-assist]

An unconditional 10-second sleep (await asyncio.sleep(10)) in the finally block of generate introduces a significant and unnecessary latency bottleneck to the rollout generation loop. Since generate is called per-sample, this sleep will block the completion of every rollout task by 10 seconds, even when the session is successfully dropped. If this delay is meant to allow background cleanup or connection teardown, it should be handled asynchronously or the timeout should be reduced/configurable.

Suggested change

	finally:
	# Close the sid before next train step's release_memory_occupation;
	# stragglers from this trajectory would otherwise race its idle assert.
	await state.adapter.finish_session(session_id) # idempotent
	await state.adapter.drop_session(session_id, wait_timeout=30) # cleanup only, idempotent
	await asyncio.sleep(10)
	finally:
	await state.adapter.drop_session(session_id, wait_timeout=30) # cleanup only, idempotent

[gemini-code-assist]

The _chat_once function performs a single HTTP POST request to the vLLM server without any retry mechanism. During long-context evaluations (which can take hours), transient network hiccups, temporary server overload, or vLLM timeouts can cause individual requests to fail, resulting in the entire sample being skipped and lost. Implementing a simple retry loop with exponential backoff would make the evaluation process significantly more robust.

async def _chat_once(
    session: aiohttp.ClientSession,
    model: str,
    messages: list[dict],
    temperature: float,
    top_p: float,
    max_tokens: int,
    retries: int = 3,
    backoff: float = 2.0,
) -> str:
    payload = dict(
        model=model,
        messages=messages,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens,
    )
    for attempt in range(retries):
        try:
            async with session.post(
                f"{BASE_URL}/chat/completions",
                headers={"Authorization": f"Bearer {API_KEY}"},
                json=payload,
            ) as resp:
                if resp.status != 200:
                    body = await resp.text()
                    raise RuntimeError(f"HTTP {resp.status}: {body[:300]}")
                data = await resp.json()
            return data["choices"][0]["message"]["content"]
        except Exception as e:
            if attempt == retries - 1:
                raise
            await asyncio.sleep(backoff * (2**attempt))

[gemini-code-assist]

The sed -i command used to patch transformer_engine is fragile because it can silently fail to perform any replacement if the target string changes or differs slightly in formatting, while still returning exit code 0. To make the build fail-fast and prevent broken images from being produced, you should use grep to verify that the replacement actually succeeded.

RUN F=$(find /usr/local/lib/python3.12/dist-packages/ -path "*transformer_engine*/pytorch/attention/dot_product_attention/utils.py" | head -1) && \
    if [ -z "$F" ]; then echo "ERROR: TE utils.py not found" && exit 1; fi && \
    sed -i 's/max_version = PkgVersion("2.8.3")/max_version = PkgVersion("2.8.4")/' "$F" && \
    grep -q 'max_version = PkgVersion("2.8.4")' "$F"

[gemini-code-assist]

Hardcoding MAX_JOBS=96 for the FlashAttention-3 compilation can easily cause out-of-memory (OOM) errors on build machines with limited RAM, as CUDA compilation is extremely memory-intensive. It is safer and more consistent to use the FA2_MAX_JOBS build argument (or a similar configurable limit) to control the compilation concurrency for both FlashAttention-2 and FlashAttention-3.

RUN git clone https://github.com/Dao-AILab/flash-attention.git && \
    cd flash-attention/ && git checkout 002cce0a1068f8c07dfccb5a1d232b9a3276947c && git submodule update --init && \
    cd hopper/ && \
    FLASH_ATTENTION_FORCE_BUILD=TRUE MAX_JOBS=${FA2_MAX_JOBS} pip -v install . --no-build-isolation && \
    cd /root/ && rm -rf flash-attention/

[gemini-code-assist]

Since tau-bench is a tool-use environment, the conversation history will contain "tool" messages. If the chat template formats "tool" messages in a way that rewrites history (similar to "user" messages in some reasoning templates), curr.startswith(prev) will be False, and the code will raise a ValueError and crash. To make the token delta calculation robust for tool-use agents, you should handle "tool" (and other potential roles) similarly to "user" by rendering them independently when they are the latest message.

Suggested change

	if curr.startswith(prev):
	new_text = curr[len(prev) :]
	elif messages[-1]["role"] == "user":
	# Reasoning templates such as Qwen3 can rewrite history when a new user
	# message arrives. Render that message independently instead of slicing
	# the rewritten conversation at the old conversation length.
	new_text = tokenizer.apply_chat_template(
	[messages[-1]],
	add_generation_prompt=False,
	tokenize=False,
	)
	if not curr.endswith(new_text):
	raise ValueError("The latest user message is not a standalone suffix of the rendered conversation")
	else:
	raise ValueError("The chat template rewrote history while calculating a non-user token delta")
	if curr.startswith(prev):
	new_text = curr[len(prev) :]
	elif messages[-1]["role"] in ("user", "tool"):
	# Reasoning and tool-use templates can rewrite history when a new message
	# arrives. Render that message independently instead of slicing
	# the rewritten conversation at the old conversation length.
	new_text = tokenizer.apply_chat_template(
	[messages[-1]],
	add_generation_prompt=False,
	tokenize=False,
	)
	if not curr.endswith(new_text):
	raise ValueError(f"The latest {messages[-1]['role']} message is not a standalone suffix of the rendered conversation")
	else:
	raise ValueError("The chat template rewrote history while calculating a non-user/non-tool token delta")

[gemini-code-assist]

There is a significant inconsistency in dependency version pinning between docker/Dockerfile and docker/Dockerfile.rocm. docker/Dockerfile pins numpy==1.26.4 and scipy==1.17.1, whereas docker/Dockerfile.rocm uses looser constraints "numpy<2" "scipy<1.16". To prevent subtle runtime discrepancies or bugs across different hardware platforms, it is highly recommended to keep dependency versions consistent.

RUN --mount=type=cache,target=/root/.cache/pip 
    pip install -r /tmp/vime-requirements.txt && 
    pip install "numpy==1.26.4" "scipy==1.17.1"