Enabling Dissag Infra for weightfree + TODO addressed

[quic-amitraj]

No description provided.

[ochougul]

Include moe_prefill_num_pipeline_stages and
moe_prefill_num_parallelized_experts here. Two expert_parallel
exports of the same model and dtype but different P/E values
produce different packed tensor shapes, but this key sends
both to the same prepared directory and manifest, so the
second export can silently reuse incompatible weights.

[amarquic]

done

[ochougul]

Add the expert-parallel transform IDs here, even if
they resolve to the unconfigured base classes for key lookup.
The manifest can now record moe_expert_parallel_stacking_v1 or
gptoss_mxfp4_dequant_expert_parallel_v1, but this lookup
cannot resolve either, so
promote_initializers_and_build_spec() loses the active
transform and per-expert .mlp. ONNX names no longer map
to .block_sparse_moe. checkpoint keys.

[amarquic]

done.

[ochougul]

why is this removed?

[amarquic]

This condition doesn't make sense for prefill moe if you see this hardcoded the number of scatter node to 2 but for prefill with expert parallel, we have more than 2 operator ctx_nodes then it will not apply this onnx transform and which raise compiler error. Now even if we remove this condition below
scatter_nodes.sort(key=cls._scatter_sort_key) # Only the first two scatter outputs map to key / value respectively. scatter_outputs = [n.output[0] for n in scatter_nodes[:2]]

we sort and select top 2 which is required for the retain_state renaming which work perfectly fine. So the condition was redundant previously.

[ochougul]

there are some models that might have single KV instead of two like deepseek there we want only first scatter output to map to KV cache for MLA. We may need to handle that when we enable deepseek models. Can you put this as a TODO here.

[ochougul]

I think this was intentional, did you check the NOTE? why remove this?

[amarquic]

revert the changes. I was changing the torch.tensor() to torch.full_like() that was the follow up changes. I reverted this .

[ochougul]

PR 1345 Checkpoint Pipeline Refactor Plan

Goal

Replace the current checkpoint-transform selector/consumed-key model with a real planning and execution pipeline for
weight-free checkpoint preparation.

The design must support:

• multiple transforms applied to the same logical tensor
• transforms that require a group of keys before they can run
• one source read path and one final write path where possible
• bounded RAM usage while still parallelizing independent work
• centralized checkpoint finalization: sidecars, index, manifest, sentinel

Current Problem

PR 1345 changes the old selector into a staged flow, but it still uses a consumed set:

for transform in transforms_to_run:
    remaining = {k: v for k, v in weight_map.items() if k not in consumed}
    result = transform.apply(src, out, target_dtype=target_dtype, weight_map=remaining)
    consumed.update(transform.get_consumed_keys(weight_map))

This means each raw key is effectively owned by one transform. That cannot model cases where the same logical tensor
needs multiple stages:

raw expert keys -> stack/dequant/split -> expert_parallel pack -> dtype/final output

The current workaround is to create special combined transform subclasses, but that does not scale and does not fully
address the pipeline TODO.

Target Model

Use a two-phase design:

1. Planning phase

   • inspect weight_map, model config, hash_params, and target dtype
   • build explicit checkpoint tasks
   • identify dependencies between tasks
   • estimate peak RAM per task
   • derive the prepared-checkpoint cache key from the full plan

2. Execution phase

   • run ready tasks through a shared thread pool
   • submit tasks only when estimated RAM budget allows
   • load required tensors once for the task
   • run all compatible transform stages in memory
   • write final tensors once
   • release tensors immediately
   • collect final weight_map

Core Types

Add small internal planning types, probably in QEfficient/base/checkpoint_transforms.py or a nearby internal module.

from dataclasses import dataclass


@dataclass(frozen=True)
class TensorRef:
    key: str
    stage: str


@dataclass(frozen=True)
class TaskParams:
    transform_id: str
    values: tuple[tuple[str, object], ...]


@dataclass
class CheckpointTask:
    task_id: str
    input_refs: tuple[TensorRef, ...]
    output_refs: tuple[TensorRef, ...]
    source_files: tuple[str, ...]
    output_file: str
    estimated_peak_bytes: int
    params: TaskParams

    def run(self, src: Path, out: Path, target_dtype: torch.dtype) -> dict[str, str]:
        ...

TensorRef.stage is important. It distinguishes:

("model.layers.0.moe_weights.gate", "stacked")
("model.layers.0.moe_weights.gate", "packed")
("model.layers.0.moe_weights.gate", "final")

Without staged tensor refs, a plain key-to-transform map cannot represent multiple transforms on one logical key.

Transform Contract

Replace or supplement the current apply() / get_consumed_keys() contract with planning:

class BaseCheckpointTransform:
    TRANSFORM_ID: str = ""

    @classmethod
    def plan_tasks(
        cls,
        weight_map: dict[str, str],
        config,
        hash_params: dict,
        target_dtype: torch.dtype,
    ) -> list[CheckpointTask]:
        ...

Transforms should declare task groups. The pipeline should own scheduling, loading, writing, and finalization.

Task Shapes

Dense / Base Dtype Conversion

Use independent shard-level or key-group tasks:

input:  raw base keys from one shard
output: final same keys
run:    load -> dtype convert if needed -> write base shard

These tasks can run immediately and in parallel, subject to RAM budget.

Per-Expert MoE Stacking

Use one grouped task per layer:

input:  all expert projection keys for that layer
output: moe_weights.gate, moe_weights.up, moe_weights.down
run:    load all required expert tensors -> stack -> optional pack -> dtype/final -> write one layer shard

The task must validate that all required experts/projections are present before execution.

GPT-OSS MXFP4 Dequant + Split

Use one grouped task per layer:

input:  gate_up/down blocks, scales, optional biases
output: canonical moe_weights.*
run:    load blocks/scales/biases -> dequant -> split -> optional pack -> dtype/final -> write one layer shard

Fused Expert Split

Use one grouped task per layer or per fused tensor prefix:

input:  gate_up_proj, down_proj, optional biases
output: canonical moe_weights.*
run:    load fused tensors -> split -> optional pack -> dtype/final -> write final shard

Expert-Parallel Packing

Do not write this as a separate intermediate disk stage.

When hash_params["moe_prefill_flavour"] == "expert_parallel", fold packing into the grouped task execution after
stack/dequant/split and before final write:

raw group -> canonical moe_weights -> expert_parallel packed moe_weights -> final write

This avoids:

• writing temporary stacked shards
• reading those shards again for packing
• growing disk usage
• creating cleanup/staleness problems

Memory-Bounded Execution

Add a scheduler that accepts planned tasks and a memory budget.

Default budget:

max_ram_bytes = int(available_ram_gb() * 0.8 * 1024**3)

Allow an internal override later if needed, but keep the first implementation simple.

Scheduler behavior:

pending = list(tasks)
running = {}
available_refs = set(raw_refs)
active_bytes = 0

while pending or running:
    for task in ready_tasks(pending, available_refs):
        if active_bytes + task.estimated_peak_bytes <= max_ram_bytes:
            future = executor.submit(task.run, src, out, target_dtype)
            running[future] = task
            active_bytes += task.estimated_peak_bytes
            pending.remove(task)

    completed_future = wait_for_one_future(running)
    task = running.pop(completed_future)
    active_bytes -= task.estimated_peak_bytes
    final_weight_map.update(completed_future.result())
    available_refs.update(task.output_refs)

If one task exceeds the cap, fail early:

raise ValueError(
    f"Checkpoint transform task {task.task_id} requires at least {required_gb:.2f} GB, "
    f"but max checkpoint transform RAM is {limit_gb:.2f} GB."
)

The cap must be allowed to be smaller than total model size, but not smaller than the largest indivisible grouped task.

Intermediate Tensor Policy

Do not write intermediate tensors by default.

Within a task:

load raw tensors
run all applicable transform stages in memory
write final output shard
release tensors

Only introduce spill-to-disk later if a real model needs it and the failure mode is understood.

Cache Key

The prepared checkpoint hash should be derived from the plan, not only from loose top-level fields.

Include at least:

• source model ref or resolved source path
• source checkpoint fingerprint
• target dtype
• task transform IDs
• task parameters
• input keys
• output keys
• output file naming
• expert-parallel parameters:
  • moe_prefill_flavour
  • moe_prefill_num_pipeline_stages
  • moe_prefill_num_parallelized_experts
  • moe_prefill_expert_parallel_chunk_size if it affects layout/output

This fixes the incomplete TODO-6 behavior where two expert_parallel layouts can collide in the same prepared directory.

Pipeline Responsibilities

CheckpointTransformPipeline.apply() should own:

• source format validation
• reading the original weight_map
• building the plan
• checking cache manifest
• clearing stale output
• invoking the memory-bounded executor
• writing final model.safetensors.index.json
• copying sidecars
• writing manifest
• touching sentinel

Transform classes should not independently write index, manifest, sidecars, or sentinel.

Manifest

Bump or extend the manifest to include:

• manifest version
• source path
• target dtype
• source file fingerprint
• plan fingerprint
• active transform IDs
• transform parameters
• final output files

The manifest should be strict enough that a changed packing layout, transform parameter, or source file invalidates the
prepared directory.

Migration Steps

1. Add TensorRef, CheckpointTask, and planning helpers.
2. Add a memory-bounded task executor.
3. Implement dense dtype conversion as planned base shard tasks.
4. Convert per-expert MoE stacking to layer-level grouped tasks.
5. Convert GPT-OSS MXFP4 dequant/split to layer-level grouped tasks.
6. Convert fused expert split to grouped tasks using canonical key mapping.
7. Fold expert-parallel packing into the relevant grouped tasks instead of using separate combined subclasses.
8. Move all finalization to CheckpointTransformPipeline.apply().
9. Derive prepared checkpoint hash from the plan fingerprint.
10. Remove or update the stale selector TODO once the new pipeline supports composable stages.

Validation

Use existing weight-free tests where possible. Add focused cases under existing test files rather than creating broad new
test files.

Minimum tests:

• dense checkpoint still performs dtype conversion and writes a valid index
• per-expert MoE stacking writes canonical final keys
• GPT-OSS MXFP4 dequant/split writes canonical final keys
• fused expert split handles Mixtral-style and GraniteMoE-style keys
• expert_parallel with different P/E values produces different prepared hashes
• one logical key can pass through multiple stages before final write
• grouped task waits until all required keys are present
• missing required grouped key fails before writing a partial prepared checkpoint
• memory budget limits concurrent grouped tasks
• one task above memory cap raises a clear error

Acceptance Criteria

• No transform owns global finalization.
• No intermediate checkpoint files are written for normal stack/dequant/split/pack flows.
• Final output tensors are written once.
• Final weight_map contains only valid final tensor keys and existing shard files.
• Cache identity changes when transform layout or expert_parallel packing parameters change.
• The pipeline can represent multiple transforms on one logical key without hardcoded combined subclasses.

[quic-amitraj]

CI-Ready

[quic-rishinr]

@quic-amitraj please add proper description

[quic-rishinr]

CI-Ready

[ochougul]

This still writes stacked experts-layer-* shards and then reopens/overwrites them for expert-parallel packing. The requested shape was raw group -> stack/dequant/split -> pack -> final write inside
one task, with no normal intermediate checkpoint stage.

[ochougul]

This is still the selector/consumed-key model from the previous review. A raw key can only be owned by one transform, so multi-stage flows still have to be encoded as combined transform classes
instead of planned tasks with staged refs. Please move this to the planned task pipeline or stop claiming the pipeline TODO is addressed.