Tracking issue for a cleanup that came up while reviewing #8185. Filed short so it's not lost; not asking for it in that PR.
Problem
tp_shard.py keeps num_kv_heads, tp_grain_size, num_attention_heads, n_embd as process-wide mutable globals (set_num_kv_heads / set_tp_grain_size during AutoTP replacement, read back in get_shard_size / get_shard_size_list).
Because they are process-wide, a second AutoTP model loaded into the same process overwrites them, so the first model's later sharding/gather silently reads the wrong values. This makes it unsafe to run more than one AutoTP model per process — teacher/student setups with different TP sizes, online distillation, or any flow that builds a second AutoTP model (including via init_inference) after the first.
#8185 patches one symptom with _freeze_partition_sizes (resolve the split once at construction, don't re-read the globals at runtime). But the globals remain the channel by which a new model's replacement clobbers the previous one's state. Removing them is the source-level fix.
Direction
Move that state off module globals onto per-model state — compute once from the model in AutoTP.__init__, thread explicitly through the sharding helpers, drop set_* / get_* / the globals. Then each model carries its own kv-head/grain state and cannot stomp on another's; multiple AutoTP models can coexist in one process. Mechanical change, many call sites, hence tracked separately.
Related
Tracking issue for a cleanup that came up while reviewing #8185. Filed short so it's not lost; not asking for it in that PR.
Problem
tp_shard.pykeepsnum_kv_heads,tp_grain_size,num_attention_heads,n_embdas process-wide mutable globals (set_num_kv_heads/set_tp_grain_sizeduring AutoTP replacement, read back inget_shard_size/get_shard_size_list).Because they are process-wide, a second AutoTP model loaded into the same process overwrites them, so the first model's later sharding/gather silently reads the wrong values. This makes it unsafe to run more than one AutoTP model per process — teacher/student setups with different TP sizes, online distillation, or any flow that builds a second AutoTP model (including via
init_inference) after the first.#8185 patches one symptom with
_freeze_partition_sizes(resolve the split once at construction, don't re-read the globals at runtime). But the globals remain the channel by which a new model's replacement clobbers the previous one's state. Removing them is the source-level fix.Direction
Move that state off module globals onto per-model state — compute once from the model in
AutoTP.__init__, thread explicitly through the sharding helpers, dropset_*/get_*/ the globals. Then each model carries its own kv-head/grain state and cannot stomp on another's; multiple AutoTP models can coexist in one process. Mechanical change, many call sites, hence tracked separately.Related
_freeze_partition_sizesthis generalizes.