xai-org/x-algorithm · error · ValueError

Cannot remap XLA compile options device assignment: assignme

Error message

Cannot remap XLA compile options device assignment: assignment expects {expected_devices} devices ({replica_count=} x {computation_count=}), but got {len(execution_devices)} execution devices.

What it means

When loading a serialized AOT-compiled executable onto a new set of devices, the stored XLA device assignment (replica_count x computation_count) must match the number of execution devices supplied. This ValueError is raised when remapping the compile options' device assignment because the device topology changed between compile time and load time.

Source

Thrown at phoenix/xrex/utils/aot.py:115

    return compile_options


def remap_compile_options_device_assignment(
    compile_options: xc.CompileOptions | None,
    execution_devices: Sequence[xc.Device],
) -> xc.CompileOptions | None:
    if compile_options is None:
        return None

    device_assignment = compile_options.device_assignment
    if device_assignment is None:
        return compile_options

    replica_count = device_assignment.replica_count()
    computation_count = device_assignment.computation_count()
    expected_devices = replica_count * computation_count
    if len(execution_devices) != expected_devices:
        raise ValueError(
            "Cannot remap XLA compile options device assignment: "
            f"assignment expects {expected_devices} devices "
            f"({replica_count=} x {computation_count=}), "
            f"but got {len(execution_devices)} execution devices."
        )

    current_device_ids = np.asarray([d.id for d in execution_devices], dtype=np.int64)
    compile_options.device_assignment = xc.DeviceAssignment.create(
        current_device_ids.reshape(replica_count, computation_count, order="F")
    )
    return compile_options


def load_compiled_from_serialized(
    lowered: Lowered,
    partially_serialized: PartiallySerialized,
    execution_devices: Sequence[xc.Device],
):

View on GitHub (pinned to 24c60942c5)

Solutions

1. Recompile (delete/invalidate the AOT cache entry) on the new topology so the device assignment matches
2. Ensure the same number of visible devices at load time as at compile time (e.g. same CUDA_VISIBLE_DEVICES)
3. Point aot_cache_dir at a per-topology cache directory keyed by device count

Example fix

# before
compiled = load_compiled_from_serialized(cache_dir, devices=devices_4)  # cache built for 8
# after
compiled = compile_or_load(..., aot_cache_dir=cache_dir_for_this_topology)  # or clear stale cache

n = jax.device_count()
# only reuse a cache compiled for the same device count
cache_key = f"{aot_cache_dir}/devices_{n}"
if not (Path(cache_key) / 'meta.json').exists():
    aot_cache_dir = None  # force fresh compile

try:
    exe = load_compiled_from_serialized(path, devices)
except ValueError as e:
    if 'execution devices' in str(e):
        exe = compile_fresh(devices)  # fallback recompile
    else:
        raise

When it happens

Trigger: Calling load_compiled_from_serialized or persistent_load with an execution_devices list whose length differs from replica_count * computation_count of the compiled artifact, e.g. loading a cache compiled for 8 devices onto 4 devices.

Common situations: Sharing an AOT cache directory across machines with different GPU counts; changing JAX device count (CUDA_VISIBLE_DEVICES) between compile and load; single-host cache reused on multi-host jobs.

Related errors

• AOT is not implemented for non-XLA lowerings.
• Compilation does not support serialization
• AOT cache directory must be provided.
• Cannot lower a compiled jit function.
• Unsupported type: {type(self.jitted)}.

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/f2f2aa26eb1c954e. Report an issue: GitHub.