runtime_utils

Utilities for runtime benchmarking and model saving in Puzzletron.

This module provides classes and utility functions used for empirical runtime estimation of Transformer subblocks and for saving models and tokenizers in formats suitable for benchmarking with vLLM.

Classes

RuntimeConfig

Configuration for a vLLM latency benchmark run.

Functions

save_model

Save model weights as AnyModel and copy the tokenizer to output_path.

save_model_as_anymodel

Save a temporary vLLM-compatible AnyModel benchmark checkpoint.
class RuntimeConfig

Bases: object

Configuration for a vLLM latency benchmark run.

__init__(vocab_size, hidden_size, num_attention_heads, num_key_value_heads, descriptor, model_config_fields, tokenizer_path, repeat_block_n_times, prefill_seq_len, generation_seq_len, batch_size, num_iters, num_warmup_iters, extra_vllm_args=(), max_num_seqs=None, topology=RuntimeTopology(tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, prefill_context_parallel_size=1, decode_context_parallel_size=1, enable_expert_parallel=False, distributed_executor_backend='mp', gpu_group_size=1), estimator_schema='candidate_slope_v1', estimator_mode='homogeneous', effective_repeat_count=None, scaffold_policy='none', vllm_env=())
Parameters:

  • vocab_size (int)

  • hidden_size (int)

  • num_attention_heads (int)

  • num_key_value_heads (int)

  • descriptor (type)

  • model_config_fields (tuple[tuple[str, Any], ...])

  • tokenizer_path (str)

  • repeat_block_n_times (int)

  • prefill_seq_len (int)

  • generation_seq_len (int)

  • batch_size (int)

  • num_iters (int)

  • num_warmup_iters (int)

  • extra_vllm_args (tuple[str, ...])

  • max_num_seqs (int | None)

  • topology (RuntimeTopology)

  • estimator_schema (str)

  • estimator_mode (str)

  • effective_repeat_count (int | None)

  • scaffold_policy (str)

  • vllm_env (tuple[tuple[str, str], ...])

Return type:

None

batch_size: int
descriptor: type
effective_repeat_count: int | None = None
estimator_mode: str = 'homogeneous'
estimator_schema: str = 'candidate_slope_v1'
extra_vllm_args: tuple[str, ...] = ()
generation_seq_len: int
hidden_size: int
max_num_seqs: int | None = None
model_config_fields: tuple[tuple[str, Any], ...]
model_config_value(key, default=None)

Return a descriptor-specific benchmark config value.

Parameters:

  • key (str)

  • default (Any)

Return type:

Any

num_attention_heads: int
num_iters: int
num_key_value_heads: int
num_warmup_iters: int
prefill_seq_len: int
repeat_block_n_times: int
scaffold_policy: str = 'none'
tokenizer_path: str
topology: RuntimeTopology = RuntimeTopology(tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, prefill_context_parallel_size=1, decode_context_parallel_size=1, enable_expert_parallel=False, distributed_executor_backend='mp', gpu_group_size=1)
vllm_env: tuple[tuple[str, str], ...] = ()
vocab_size: int
save_model(model, tokenizer_path, output_path, descriptor)

Save model weights as AnyModel and copy the tokenizer to output_path.

Parameters:

  • model (PreTrainedModel)

  • tokenizer_path (Path)

  • output_path (Path)

  • descriptor (type)

Return type:

None

save_model_as_anymodel(model, output_dir, descriptor, runtime_descriptor=None)

Save a temporary vLLM-compatible AnyModel benchmark checkpoint.

Parameters:

output_dir (Path)