[FlexCheckpoint] Add LoadTransform extension point for quantized checkpoints - #79712
[FlexCheckpoint] Add LoadTransform extension point for quantized checkpoints#79712Lmywl wants to merge 2 commits into
Conversation
…kpoints FlexCheckpoint could only load checkpoints whose physical tensors map one-to-one onto logical model weights. Quantized checkpoints break that assumption: one logical weight is stored as several physical tensors (packed values plus scales), often in element formats that have no exact Paddle dtype. This adds `paddle.distributed.LoadTransform`, a format-independent protocol that lets a caller declare virtual logical tensors, list the physical checkpoint tensors each one needs, and materialize the logical tensor once those components have been assembled. The transform runs after resharding, so dequantization stays out of the reshard and communication layers. Supporting changes: - `load_state_dict` / `load_state_dict_impl` accept `load_transform` and split the load dict into passthrough targets and physical components, rejecting key collisions between the two. - AOA sees logical tensors instead of the physical components they consume, so renames and other statements apply to the logical key. - Safetensors one-byte formats (F8_E4M3, F8_E4M3FN, F8_E8M0) are transported losslessly as uint8; unregistered formats now raise instead of being silently reinterpreted. - `create_hf_ckpt_metadata` returns the metadata it builds and no longer assumes a distributed environment is initialized. - Tuple type hints in `metadata.py` are made variadic.
| ) | ||
| if piece.place != target_tensor.place: | ||
| piece = piece.to(target_tensor.place) | ||
| paddle.assign(piece, target_tensor) |
There was a problem hiding this comment.
在 AOA + send_recv/grouped_send_recv 路径中，前面会对需要临时 logical shard 的目标调用 _clear_to_zero_allocation()；这里直接 paddle.assign(piece, target_tensor)，却没有先为未初始化 Tensor 分配并共享 buffer。此时 transform 应用会在 assign_out_ 处失败，量化 checkpoint 无法完成。请像现有 reshard 通信路径一样，在赋值前对 _is_initialized() == False 的目标用 zeros_like(...)._share_buffer_to(...) 初始化，并覆盖 AOA 生成的临时目标。
There was a problem hiding this comment.
Thanks — I tried to reproduce this and could not, so I have left the code unchanged. Please push back if you have a case that does fail.
paddle.assign() allocates its output, so a whole-tensor assign into a _clear_to_zero_allocation()'d target succeeds (paddle 3.4.0, gpu:0):
before: _is_initialized = False
after : _is_initialized = True shape = [4, 8] correct = True
sliced: _is_initialized = True shape = [2, 8] correct = True # source is a paddle.slice() view
cross : _is_initialized = True place = gpu:0 correct = True # source on another place
What the zeros_like(...)._share_buffer_to(...) guard in resharder.py protects against is a sub-region write, which cannot allocate because the written region is smaller than the tensor:
t = paddle.zeros([4, 8]); t._clear_to_zero_allocation()
t[0:2] = src[0:2]
ValueError: (InvalidArgument) The size of Holder is not enough to store the Tensor.
The transform apply path only ever does a whole-tensor assign of a piece that has exactly the target's local shape, so it is the first case, not the second. Adding the guard here would add a full-size zeros_like allocation plus a zero-fill that the immediately following whole-tensor assign makes redundant, on the path where memory is most scarce.
On AOA temporary targets specifically: they keep being written after the fix for the sharded-target finding below. They are ShardedWeights, and extract_tensor_metadata() reports None only for an uninitialized bare paddle.Tensor, i.e. the documented "this tensor is placed in another mesh which does not contain this rank" case.
There was a problem hiding this comment.
感谢补充实测。结合仓库中的 assign_out_ 实现和当前代码对目标 local shape 的整块赋值，确认此前关于未初始化目标会在 assign_out_ 失败的判断不成立；此项无需代码修改，接受该说明。
| global_offset=tuple(target.global_offset), | ||
| dtype=str(target.local_tensor.dtype).removeprefix('paddle.'), | ||
| ) | ||
| shape = tuple(target.shape) |
There was a problem hiding this comment.
当调用方按文档传入 dist.shard_tensor(...) 目标（而不是 ShardedWeight）时，这里把 target.shape 当作 local_shape、把 offset 固定为 0；但分布式 Tensor 的 shape 是全局 shape。_build_transform_component_load_dict 因而向 read_plan 传入错误的 shard 元数据，导致物理组件读取/切片错误；应用阶段也必须写入 target._local_value()，不能把结果直接赋给分布式包装器。请统一计算分布式目标的 local shape/global offset，并补充多卡 sharded target 的 transform 回归测试。
There was a problem hiding this comment.
Agreed, fixed in 609df8e.
_target_shard_metadata() is gone. Every target now goes through extract_tensor_metadata() (dcp/utils.py), the helper the rest of the load path already uses, so the local shape and global offset come from compute_local_shape_and_global_offset(shape, process_mesh, placements) and the result is written into _local_value(). A target that is not initialized on this rank is skipped instead of being described by its global shape.
One correction to the diagnosis, for the record: the values were not wrong in the simple case. paddle.assign() into a distributed tensor reshards, so writing a whole (4, 8) tensor into a Shard(0) target still left each rank holding its own rows. What was actually broken is that read_plan() was handed the global shape — so a transform that narrows its read got a plan covering the whole tensor and every rank read the full physical tensors, which defeats the purpose of read_plan() for a quantized weight — and that targets living on a mesh without this rank were written anyway.
Regression test: test/flex_checkpoint/test_load_transform_dist.py, two cards, loading into dist.shard_tensor(..., [dist.Shard(0)]). The global case (transform without read_plan()) passes before and after, for the reason above. The local case fails before the fix with AssertionError: rank 0 was told local shape (4, 8) and passes after. Single-process coverage of the target description itself is in TestLoadTransformComponentPlan (plain tensor, ShardedWeight, and a target absent on the rank).
There was a problem hiding this comment.
已核对 609df8e8：目标统一通过 extract_tensor_metadata() 获取局部 shape/offset，并写入分布式 Tensor 的 local value；新增双卡 global/local 用例覆盖该路径。此项已解决。
| } | ||
| for logical_key in sorted(targets_by_logical_key): | ||
| logical_metadata = transform_metadata[logical_key] | ||
| read_plan = getattr(load_transform, "read_plan_for", lambda _key: None)( |
There was a problem hiding this comment.
协议把 read_plan() 和 read_plan_for() 都描述为可选，但这里仅在构建阶段调用 read_plan；应用阶段若未实现 read_plan_for 就默认按 global 处理。于是 read_plan 返回 local 时，physical source 已是本地 shard，apply 的输出又按全局 offset 被切一次（例如 offset 为 (2, 0) 时会得到错误或越界数据）。请复用构建阶段的 plan，或在 local plan 缺少对应应用信息时明确报错，并增加只实现 read_plan 的回归测试。
There was a problem hiding this comment.
Agreed, fixed in 609df8e.
_build_transform_component_load_dict() now returns (components, plans) and _apply_load_transform() takes those plans, so the apply phase reuses the plan the components were actually read with instead of deciding again. read_plan_for() is removed from the protocol, and the docstring now states the contract explicitly: a "local" plan means apply() receives local source shards and must return the local shard of the logical tensor.
One note on the symptom, since it affects how urgent this looks: it was not out-of-range data but a silent no-write. paddle.slice clamps out-of-range ends and paddle.assign resizes its output, so with global offset (2, 0) the second slice yields shape [0, 8] and the target is silently resized to [0, 8]:
slice of a local shard by the target's global offset -> [0, 8]
target after assign -> [0, 8] | numel = 0
Regression tests: TestApplyLoadTransform.test_reuses_the_local_plan_the_components_were_built_from (a transform that implements only read_plan()), plus the local case of the new two-card test_load_transform_dist.py.
There was a problem hiding this comment.
已核对 609df8e8：构建阶段返回的 read plan 现在直接传给应用阶段，协议也移除 read_plan_for() 依赖；新增仅实现 read_plan() 的回归测试。此项已解决。
There was a problem hiding this comment.
Paddle-Bot Review Board （review完成)
| 序号 | 位置 | 优先级 | 规则来源 | 状态 |
|---|---|---|---|---|
| 1 | AOA 临时目标写回 |  |
仓库规则：分布式与资源初始化 | 🟡 |
| 2 | 分布式目标元数据 | |
仓库规则：分布式目标与局部 shape | ✅ |
| 3 | local read plan 一致性 | |
仓库规则：跨模块契约与分布式读取 | ✅ |
Failed CI看板
日志分析报告失败的测试 case: 根本原因分析: REST 分页筛选当前 PR head Linux-XPU 失败与 PR 新增测试直接相关。 CI-H 的 Job 状态确认为失败，但逐 Job 日志抓取失败，因此没有足够证据判断是 PR 代码、NCCL、测试机还是其他基础设施问题。 修复建议:
Powered by Nyanpasu with gpt-5.6-luna 默认推理级别, please check the suggestions carefully. |
A distributed target reports its global shape, so deriving a transform target's local shape and global offset from `shape` described the wrong region: `read_plan()` was handed the global shape and every rank read whole physical tensors, and a target placed on a mesh that does not contain the current rank was written anyway. Route every target through `extract_tensor_metadata`, the helper the rest of the load path already uses, which computes the shard from the target's placements, writes into `_local_value()`, and reports `None` for a target this rank does not hold. The build and apply phases also decided local-versus-global reads independently: components were built from `read_plan()` while the write back consulted `read_plan_for()`. A transform providing only the former therefore had its already-local output sliced a second time by the target's global offset, which silently yields an empty tensor rather than raising. The build phase now hands its plans to the apply phase, and the redundant `read_plan_for()` entry point is gone. Tests: `test_load_transform` covers both crosses plus targets absent on a rank; the new two-card `test_load_transform_dist` loads into a `Shard(0)` target, with and without `read_plan()`.
|
/re-run all-failed |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #79712 +/- ##
===========================================
Coverage ? 100.00%
===========================================
Files ? 5
Lines ? 154
Branches ? 0
===========================================
Hits ? 154
Misses ? 0
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/re-run all-failed |
3 participants
PR Category
Auto Parallel
PR Types
New features
Description
FlexCheckpoint currently assumes every physical checkpoint tensor maps one-to-one onto a logical model weight. Quantized checkpoints break that assumption: a single logical weight is stored as several physical tensors (packed values plus scales), frequently in element formats that have no exact Paddle dtype.
This PR introduces
paddle.distributed.LoadTransform, a format-independent protocol that lets the caller own the numerical interpretation while FlexCheckpoint keeps owning I/O, resharding, and AOA.A transform declares three things:
logical_metadata()— the virtual logical tensors it exposessource_keys(logical_key)— the physical checkpoint tensors each logical tensor consumesapply(logical_key, source_tensors, output_dtype)— how to materialize the logical tensorA transform may additionally provide
read_plan(logical_key, target_metadata, force_global=False)to narrow the physical read to the shard a rank actually needs;target_metadatadescribes the local shard of the target on that rank. When the returned plan is"local",apply()receives local source shards and returns the local shard of the logical tensor; otherwise it receives and returns whole tensors.The transform runs after the physical components have been fully assembled, so dequantization logic stays out of the reshard and communication layers.
Supporting changes
load_state_dict/load_state_dict_impltake aload_transformargument and split the load dict into passthrough targets and physical components. Collisions between component keys and passthrough targets are rejected explicitly.F8_E4M3,F8_E4M3FN,F8_E8M0) are transported losslessly asuint8. Previouslyassert dtype in dtype_mappingrejected them; unregistered formats now raise a descriptiveValueErrorrather than being silently reinterpreted.create_hf_ckpt_metadatareturns the metadata it builds, and no longer assumes a distributed environment is initialized (is_initialized()guard beforeget_world_size()).metadata.py(tuple[int]->tuple[int, ...]).Tests
test/flex_checkpoint/test_load_transform.py, 31 single-process cases covering component planning, offload placement, local vs. global read plans, target shard description (plain tensor,ShardedWeight, and a target absent on the rank), dtype validation, AOA interaction with logical keys, master-weight renaming, FP8 safetensors transport losslessness, dtype-alias restoration (including on failure), and end-to-endload_state_dictwith a dequantizing transform.test/flex_checkpoint/test_load_transform_dist.py, a two-card case that loads into adist.shard_tensor(..., [dist.Shard(0)])target, once with a transform that only implements the required three methods and once with a transform that also implementsread_plan()and asserts it is told its own shard. The second case fails if a target's local shard is derived from its (global)shape.Verified passing on both CPU-only and single-GPU runs, plus the two-card suite on 2 GPUs; existing
test/flex_checkpointsingle-process suites remain green.Backward compatibility
load_transformdefaults toNoneand every new code path is gated on it, so existing load behaviour is unchanged.是否引起精度变化
否