Skip to content

Fix total_cores based on TS value #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kdulla merged 1 commit into from
Aug 6, 2026
Merged

Fix total_cores based on TS value #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions QEfficient/base/modeling_qeff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -632,6 +632,7 @@ def get_onnx_path(
num_devices: int = 1,
cores_per_expert: Optional[int] = None,
tree_reduce: Optional[bool] = None,
mdp_num_partitions: Optional[int] = 1,
**compiler_options,
):
kwargs = {
Expand All @@ -652,6 +653,7 @@ def get_onnx_path(
kwargs["num_cores"] = compiler_options.get("aic_num_cores", constants.DEFAULT_AIC_NUM_CORES)
kwargs["cores_per_expert"] = 1 if cores_per_expert is None else cores_per_expert
kwargs["tree_reduce"] = False if tree_reduce is None else tree_reduce
kwargs["mdp_num_partitions"] = 1 if mdp_num_partitions is None else mdp_num_partitions
kwargs["moe_prefill_packed_chunk_size"] = (
constants.MOE_PREFILL_PACKED_CHUNK_SIZE
if moe_prefill_packed_chunk_size is None
Expand Down Expand Up @@ -1059,6 +1061,7 @@ def _compile(
expert_parallel=expert_parallel,
cores_per_expert=cores_per_expert,
tree_reduce=tree_reduce,
mdp_num_partitions=mdp_num_partitions,
_layerwise_cache_probe=layerwise_cache_probe,
kv_cache_prefix=kv_cache_prefix,
**compiler_options,
Expand Down
36 changes: 30 additions & 6 deletions QEfficient/transformers/models/modeling_auto.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,8 @@ def _configure_vlm_moe_expert_parallel(
cores_per_expert: int,
tree_reduce: bool,
moe_prefill_packed_chunk_size: int,
prefill_only: bool = False,
mdp_num_partitions: Optional[int] = 1,
) -> None:
if num_devices <= 0:
raise ValueError("`num_devices` must be greater than 0 when configuring MoE expert parallelism.")
Expand All @@ -135,7 +137,19 @@ def _configure_vlm_moe_expert_parallel(
if not isinstance(tree_reduce, bool):
raise TypeError("`tree_reduce` must be a boolean.")

total_avl_cores = num_devices * num_cores
num_ts_devices = num_devices
if prefill_only:
if mdp_num_partitions is None:
mdp_num_partitions = 1
if mdp_num_partitions <= 0:
raise ValueError("`mdp_num_partitions` must be greater than 0 when configuring MoE expert parallelism.")
if num_devices % mdp_num_partitions != 0:
raise ValueError(
f"num_devices ({num_devices}) must be divisible by mdp_num_partitions ({mdp_num_partitions}) so "
"every prefill pipeline partition gets the same number of tensor-sliced devices."
)
num_ts_devices = num_devices // mdp_num_partitions
total_avl_cores = num_ts_devices * num_cores
if expert_parallel:
hash_params["expert_parallel"] = True
hash_params["moe_prefill_total_avl_cores"] = total_avl_cores
Expand Down Expand Up @@ -188,20 +202,20 @@ def _configure_vlm_moe_expert_parallel(
f"num_experts ({num_experts}) must be divisible by num_pipeline_stages ({num_pipeline_stages})."
)
num_parallelized_experts = num_experts // num_pipeline_stages
if num_parallelized_experts % num_devices != 0:
if num_parallelized_experts % num_ts_devices != 0:
raise ValueError(
f"num_parallelized_experts ({num_parallelized_experts}) must be divisible by num_devices "
f"({num_devices})."
f"({num_ts_devices})."
)
if num_devices > 1 and num_cores % cores_per_expert != 0:
if num_ts_devices > 1 and num_cores % cores_per_expert != 0:
raise ValueError(f"num_cores ({num_cores}) must be divisible by cores_per_expert ({cores_per_expert}).")
module.num_experts = num_experts
module.num_devices = num_devices
module.num_devices = num_ts_devices
module.cores_per_expert = cores_per_expert
module.total_avl_cores = total_avl_cores
module.num_pipeline_stages = num_pipeline_stages
module.num_parallelized_experts = num_parallelized_experts
module.experts_per_soc = num_cores // cores_per_expert if num_devices > 1 else None
module.experts_per_soc = num_cores // cores_per_expert if num_ts_devices > 1 else None
module.tree_reduce = tree_reduce
module.expert_blocking_num_nsp = num_parallelized_experts
module.expert_blocking_packed_chunk_size = moe_prefill_packed_chunk_size
Expand Down Expand Up @@ -1381,6 +1395,7 @@ def export(
cores_per_expert: int = 1,
tree_reduce: bool = False,
moe_prefill_packed_chunk_size: int = constants.MOE_PREFILL_PACKED_CHUNK_SIZE,
mdp_num_partitions: Optional[int] = 1,
kv_cache_prefix: Optional[str] = None,
**kwargs,
):
Expand Down Expand Up @@ -1428,6 +1443,8 @@ def export(
cores_per_expert=cores_per_expert,
tree_reduce=tree_reduce,
moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size,
prefill_only=bool(prefill_only),
mdp_num_partitions=mdp_num_partitions,
)

if QEfficient.base.modeling_qeff.QEFFBaseModel._layerwise_active:
Expand Down Expand Up @@ -1648,6 +1665,7 @@ def export(
cores_per_expert: int = 1,
tree_reduce: bool = False,
moe_prefill_packed_chunk_size: int = constants.MOE_PREFILL_PACKED_CHUNK_SIZE,
mdp_num_partitions: Optional[int] = 1,
layerwise: bool = False,
layerwise_window_size: int = 1,
kv_cache_prefix: Optional[str] = None,
Expand Down Expand Up @@ -1690,6 +1708,7 @@ def export(
cores_per_expert=cores_per_expert,
tree_reduce=tree_reduce,
moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size,
mdp_num_partitions=mdp_num_partitions,
layerwise_window_size=layerwise_window_size,
kv_cache_prefix=kv_cache_prefix,
**kwargs,
Expand Down Expand Up @@ -1787,6 +1806,7 @@ def export(
cores_per_expert=cores_per_expert,
tree_reduce=tree_reduce,
moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size,
mdp_num_partitions=mdp_num_partitions,
prefill_seq_len=prefill_seq_len,
_layerwise_cache_probe=layerwise_cache_probe,
kv_cache_prefix=kv_cache_prefix,
Expand Down Expand Up @@ -2199,6 +2219,7 @@ def compile(
cores_per_expert=cores_per_expert,
tree_reduce=tree_reduce,
moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size,
mdp_num_partitions=compiler_options.get("mdp_num_partitions", 1),
prefill_seq_len=prefill_seq_len,
_layerwise_cache_probe=layerwise_cache_probe,
kv_cache_prefix=kv_cache_prefix,
Expand Down Expand Up @@ -2871,6 +2892,7 @@ def export(
cores_per_expert: int = 1,
tree_reduce: bool = False,
moe_prefill_packed_chunk_size: int = constants.MOE_PREFILL_PACKED_CHUNK_SIZE,
mdp_num_partitions: Optional[int] = 1,
kv_cache_prefix: Optional[str] = None,
**kwargs,
) -> str:
Expand Down Expand Up @@ -2910,6 +2932,8 @@ def export(
cores_per_expert=cores_per_expert,
tree_reduce=tree_reduce,
moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size,
prefill_only=bool(prefill_only),
mdp_num_partitions=mdp_num_partitions,
)

inputs = self.model.get_dummy_inputs(comp_ctx_lengths=self.comp_ctx_lengths_decode)
Expand Down