aitemplate.testing
detect_target
Automatic detect target for testing
Functions:
|
Detect GPU target based on nvidia-smi and rocminfo |
benchmark_pt
helper function to benchmark eager pytorch
Functions:
|
function for benchmarking a pytorch function. |
- aitemplate.testing.benchmark_pt.benchmark_torch_function(iters: int, function, *args, **kwargs) float[source]
function for benchmarking a pytorch function.
- Parameters:
iters (int) – Number of iterations.
function (lambda function) – function to benchmark.
args (Any type) – Args to function.
- Returns:
Runtime per iteration in ms.
- Return type:
float
benchmark_ait
Functions:
|
Make input and output pools for benchmarking. |
|
Run the benchmark. |
|
Run the module with the given inputs and outputs pools. |
- aitemplate.testing.benchmark_ait.make_input_output_pools(*, pool_size, eval_pt_func, input_filter_func, output_filter_func)[source]
Make input and output pools for benchmarking. The rationale is avoiding retrieving the same input from the device cache for fair perf assessment. :param pool_size: The size of the pool. :type pool_size: int :param eval_pt_func: A callable that returns a dict of inputs and outputs. :type eval_pt_func: callable :param input_filter_func: A callable that takes a key and a value and returns True if the key-value pair from the eval_pt_func result should be included in the input pool. :type input_filter_func: callable :param output_filter_func: A callable that takes a key and a value and returns True if the key-value pair from the eval_pt_func result should be included in the output pool. :type output_filter_func: callable
- Returns:
inputs_pool (List[Dict[str, torch.Tensor]]) – A list of inputs to pass into Model.RunWithTensors.
outputs_pool (List[Dict[str, torch.Tensor]]) – A list of outputs to pass into Model.RunWithTensors.
- aitemplate.testing.benchmark_ait.run_benchmark(*, ait_module, inputs_pool, outputs_pool, num_iters, num_warmup_iters, stream: Optional[Stream] = None, sync: bool = False, graph_mode: bool = False)[source]
Run the benchmark. :param ait_module: The AIT module to run. :type ait_module: Model :param inputs_pool: A list of inputs to pass into Model.RunWithTensors. :type inputs_pool: List[Dict[str, torch.Tensor]] :param outputs_pool: A list of outputs to pass into Model.RunWithTensors. :type outputs_pool: List[Dict[str, torch.Tensor]] :param num_iters: The number of iterations to run. :type num_iters: int :param num_warmup_iters: The number of warmup iterations to run. :type num_warmup_iters: int :param stream: The CUDA stream to run the module on; if None, use the default stream. :type stream: Optional[torch.cuda.Stream] :param sync: Whether to synchronize the CUDA stream after each iteration. :type sync: bool :param graph_mode: Whether to run the module in graph mode. :type graph_mode: bool
- Returns:
The average time per iteration in milliseconds.
- Return type:
float
- aitemplate.testing.benchmark_ait.run_module_with_pools(*, ait_module, inputs_pool, outputs_pool, num_iters, stream_ptr: Optional[int] = None, sync: bool = False, graph_mode: bool = False)[source]
Run the module with the given inputs and outputs pools. :param ait_module: The AIT module to run. :type ait_module: Model :param inputs_pool: A list of inputs to pass into Model.RunWithTensors. :type inputs_pool: List[Dict[str, torch.Tensor]] :param outputs_pool: A list of outputs to pass into Model.RunWithTensors. :type outputs_pool: List[Dict[str, torch.Tensor]] :param num_iters: The number of iterations to run. :type num_iters: int :param stream_ptr: The CUDA stream pointer to run the module on; if None, use the legacy stream. :type stream_ptr: Optional[int] :param sync: Whether to synchronize the CUDA stream after each iteration. :type sync: bool :param graph_mode: Whether to run the module in graph mode. :type graph_mode: bool