usrbinenv python hf trainer benchmarking tool this tool can be used to run and compare multiple dimensions of the hf trainers args it then prints a report once in github format with all the information that needs to be shared with others and second time in a consolefriendly format so it s easier to use for tuning things up the main idea is trainerbenchmark py basecmd cmd args that don t change variations tf32 0tf32 1 fp16 0fp16 1bf16 1 targetmetrickey trainsamplespersecond the variations can be any command line argument that you want to compare and not just dtype as in the example variations allows you to compare variations in multiple dimensions as the first dimention has 2 options and the second 3 in our example this will run the trainer 6 times adding one of 1 tf32 0 fp16 0 2 tf32 0 fp16 1 3 tf32 0 bf16 1 4 tf32 1 fp16 0 5 tf32 1 fp16 1 6 tf32 1 bf16 1 and print the results this is just a cartesian product and more than 2 dimensions can be used if you want to rely on defaults this variations tf32 0tf32 1 fp16 0fp16 1bf16 1 is identical to this variations tf32 0tf32 1 fp16bf16 the leading empty variation in the 2nd dimension is a valid variation so here we get the following 6 variations 1 tf32 0 2 tf32 0 fp16 3 tf32 0 bf16 4 tf32 1 5 tf32 1 fp16 6 tf32 1 bf16 in this particular case we don t know what the default tf32 setting is as it s normally pytorchversion dependent that s why it s best to do an explicit setting of each variation tf32 0tf32 1 here is a full example of a train cudavisibledevices0 python scriptsbenchmarktrainerbenchmark py basecmd examplespytorchtranslationruntranslation py modelnameorpath t5small outputdir outputdir dotrain labelsmoothing 0 1 loggingstrategy no savestrategy no perdevicetrainbatchsize 32 maxsourcelength 512 maxtargetlength 512 numtrainepochs 1 overwriteoutputdir sourcelang en targetlang ro datasetname wmt16 datasetconfig roen sourceprefix translate english to romanian warmupsteps 50 maxtrainsamples 20000 dataloadernumworkers 2 targetmetrickey trainsamplespersecond repeattimes 1 variations fp16bf16 tf32 0tf32 1 reportmetrickeys trainloss repeattimes 1 basevariation tf32 0 and here is a possible output variation train diff train samples loss per second tf32 0 285 11 0 2 51 tf32 1 342 09 20 2 51 fp16 tf32 0 423 49 49 2 51 fp16 tf32 1 423 13 48 2 51 bf16 tf32 0 416 80 46 2 52 bf16 tf32 1 415 87 46 2 52 so you can quickly compare the different outcomes typically running each experiment once is enough but if the environment is unstable you can rerun each multiple times e g 3 using repeattimes 3 and it will report the averaged results by default it ll use the lowest result as the base line to use as 100 and then compare the rest to it as can be seen from the table above but you can also specify which combination is the one to use as the baseline e g to change to another entry use basevariation tf32 1 fp16 0 targetmetrickey is there to tell the program which metrics to compare the different metric keys are inside outputdirallresults json e g to measure eval performance instead of train use targetmetrickey evalsamplespersecond but of course you will need to adjust the basecmd value in the example to perform evaluation as well as currently it doesn t a helper class to tee print s output into a file usage sys stdout teefilename strip tqdm codes return the original command line string that can be replayed nicely and wrapped for 80 char width args maxwidth int optional defaults to 80 the width to wrap for fullpythonpath bool optional defaults to false whether to replicate the full path or just the last segment i e python deal with critical env vars python executable not always needed if the script is executable now the normal args split up into up to maxwidth lines with shell multiline escapes unwrap multiline input remove outputdir if any and set our own ensure we have overwriteoutputdir enable to debug everything but the run itself to do it fast and see the progress this is useful for debugging the output formatting quickly we can remove it later once everybody is happy with the output save the streams filter out just the keys we want def processresultsresults targetmetrickey reportmetrickeys basevariation outputdir df pd dataframeresults variationkey variation diffkey diff sentinelvalue nan if basevariation is not none and lendfdfvariationkey basevariation this may still return nan sentinelvalue df locdfvariationkey basevariationtargetmetrickey item if math isnansentinelvalue as a fallback use the minimal value as the sentinel sentinelvalue df locdftargetmetrickey nantargetmetrickey min create diff column if possible if not math isnansentinelvalue dfdiffkey df apply lambda r round100 rtargetmetrickey sentinelvalue sentinelvalue if not math isnanrtargetmetrickey else 0 axiscolumns reorder columns cols variationkey targetmetrickey diffkey reportmetrickeys df df reindexcols axiscolumns reorder cols capitalize df df renamestr capitalize axiscolumns make the cols as narrow as possible dfgithub df renamelambda c c replace br axiscolumns dfconsole df renamelambda c c replace n axiscolumns report copy between the cutherelines and paste as is to github or a forum report 88 report results dfgithub tomarkdownindexfalse floatfmt 2f report report setup getversions report the benchmark command line was getoriginalcommand report report 88 report results console dfconsole tomarkdownindexfalse floatfmt 2f printnn joinreport def main parser argparse argumentparser parser addargument basecmd defaultnone typestr requiredtrue helpbase cmd parser addargument variations defaultnone typestr nargs requiredtrue helpmultidimensional variations example fp16bf16 tf32 parser addargument basevariation defaultnone typestr helpbaseline variation to compare to if none the minimal target value will be used to compare against parser addargument targetmetrickey defaultnone typestr requiredtrue helptarget metric key in outputdirallresults json e g trainsamplespersecond parser addargument reportmetrickeys default typestr helpreport metric keys other metric keys from outputdirallresults json to report e g trainloss use a single argument e g trainloss trainsamples parser addargument repeattimes default1 typeint helphow many times to rerun each variation an average will be reported parser addargument outputdir defaultoutputbenchmark typestr helpthe output directory where all the benchmark reports will go to and additionally this directory will be used to override outputdir in the script that is being benchmarked parser addargument verbose defaultfalse actionstoretrue helpwhether to show the outputs of each run or just the benchmark progress args parser parseargs outputdir args outputdir pathoutputdir mkdirexistoktrue basecmd getbasecommandargs outputdir split each dimension into its foo variations dims listmapstr strip re splitr x for x in args variations build a cartesian product of dimensions and convert those back into cmdline arg strings while stripping white space for inputs that were empty variations listmapstr strip map join itertools productdims longestvariationlen maxlenx for x in variations split wanted keys reportmetrickeys args reportmetrickeys split capture prints into a log file for convenience reportfn fbenchmarkreportdatetime datetime now strftime ymdhms txt printfnnote each run s output is also logged under outputdirlog std txt printfand this script s output is also piped into reportfn sys stdout teereportfn printfn running lenvariations benchmarks printfbase command joinbasecmd variationkey variation results for id variation in enumeratetqdmvariations desctotal completion leavefalse cmd basecmd variation split results append processrun id 1 cmd variationkey variation longestvariationlen args targetmetrickey reportmetrickeys args repeattimes outputdir args verbose processresultsresults args targetmetrickey reportmetrickeys args basevariation outputdir if name main main usr bin env python hf trainer benchmarking tool this tool can be used to run and compare multiple dimensions of the hf trainers args it then prints a report once in github format with all the information that needs to be shared with others and second time in a console friendly format so it s easier to use for tuning things up the main idea is trainer benchmark py base cmd cmd args that don t change variations tf32 0 tf32 1 fp16 0 fp16 1 bf16 1 target metric key train_samples_per_second the variations can be any command line argument that you want to compare and not just dtype as in the example variations allows you to compare variations in multiple dimensions as the first dimention has 2 options and the second 3 in our example this will run the trainer 6 times adding one of 1 tf32 0 fp16 0 2 tf32 0 fp16 1 3 tf32 0 bf16 1 4 tf32 1 fp16 0 5 tf32 1 fp16 1 6 tf32 1 bf16 1 and print the results this is just a cartesian product and more than 2 dimensions can be used if you want to rely on defaults this variations tf32 0 tf32 1 fp16 0 fp16 1 bf16 1 is identical to this variations tf32 0 tf32 1 fp16 bf16 the leading empty variation in the 2nd dimension is a valid variation so here we get the following 6 variations 1 tf32 0 2 tf32 0 fp16 3 tf32 0 bf16 4 tf32 1 5 tf32 1 fp16 6 tf32 1 bf16 in this particular case we don t know what the default tf32 setting is as it s normally pytorch version dependent that s why it s best to do an explicit setting of each variation tf32 0 tf32 1 here is a full example of a train cuda_visible_devices 0 python scripts benchmark trainer benchmark py base cmd examples pytorch translation run_translation py model_name_or_path t5 small output_dir output_dir do_train label_smoothing 0 1 logging_strategy no save_strategy no per_device_train_batch_size 32 max_source_length 512 max_target_length 512 num_train_epochs 1 overwrite_output_dir source_lang en target_lang ro dataset_name wmt16 dataset_config ro en source_prefix translate english to romanian warmup_steps 50 max_train_samples 20000 dataloader_num_workers 2 target metric key train_samples_per_second repeat times 1 variations fp16 bf16 tf32 0 tf32 1 report metric keys train_loss repeat times 1 base variation tf32 0 and here is a possible output variation train diff train samples loss per second tf32 0 285 11 0 2 51 tf32 1 342 09 20 2 51 fp16 tf32 0 423 49 49 2 51 fp16 tf32 1 423 13 48 2 51 bf16 tf32 0 416 80 46 2 52 bf16 tf32 1 415 87 46 2 52 so you can quickly compare the different outcomes typically running each experiment once is enough but if the environment is unstable you can re run each multiple times e g 3 using repeat times 3 and it will report the averaged results by default it ll use the lowest result as the base line to use as 100 and then compare the rest to it as can be seen from the table above but you can also specify which combination is the one to use as the baseline e g to change to another entry use base variation tf32 1 fp16 0 target metric key is there to tell the program which metrics to compare the different metric keys are inside output_dir all_results json e g to measure eval performance instead of train use target metric key eval_samples_per_second but of course you will need to adjust the base cmd value in the example to perform evaluation as well as currently it doesn t a helper class to tee print s output into a file usage sys stdout tee filename strip tqdm codes return the original command line string that can be replayed nicely and wrapped for 80 char width args max_width int optional defaults to 80 the width to wrap for full_python_path bool optional defaults to false whether to replicate the full path or just the last segment i e python deal with critical env vars python executable not always needed if the script is executable now the normal args split up into up to max_width lines with shell multi line escapes unwrap multi line input remove output_dir if any and set our own ensure we have overwrite_output_dir enable to debug everything but the run itself to do it fast and see the progress this is useful for debugging the output formatting quickly we can remove it later once everybody is happy with the output save the streams filter out just the keys we want datetime datetime datetime now strftime y m d h m s software transformers transformers __version__ torch torch __version__ cuda torch version cuda python platform python_version hardware torch cuda device_count gpus properties name properties total_memory 2 30 0 2f gb this may still return nan as a fallback use the minimal value as the sentinel create diff column if possible re order columns reorder cols capitalize make the cols as narrow as possible split each dimension into its foo variations build a cartesian product of dimensions and convert those back into cmd line arg strings while stripping white space for inputs that were empty split wanted keys capture prints into a log file for convenience
import argparse
import datetime
import io
import itertools
import json
import math
import os
import platform
import re
import shlex
import subprocess
import sys
from pathlib import Path
from statistics import fmean

import pandas as pd
import torch
from tqdm import tqdm

import transformers


nan = float("nan")


class Tee:
    

    def __init__(self, filename):
        self.stdout = sys.stdout
        self.file = open(filename, "a")

    def __getattr__(self, attr):
        return getattr(self.stdout, attr)

    def write(self, msg):
        self.stdout.write(msg)
        
        self.file.write(re.sub(r"^.*\r", "", msg, 0, re.M))


def get_original_command(max_width=80, full_python_path=False):
    

    cmd = []

    
    env_keys = ["CUDA_VISIBLE_DEVICES"]
    for key in env_keys:
        val = os.environ.get(key, None)
        if val is not None:
            cmd.append(f"{key}={val}")

    
    python = sys.executable if full_python_path else sys.executable.split("/")[-1]
    cmd.append(python)

    
    cmd += list(map(shlex.quote, sys.argv))

    
    lines = []
    current_line = ""
    while len(cmd) > 0:
        current_line += f"{cmd.pop(0)} "
        if len(cmd) == 0 or len(current_line) + len(cmd[0]) + 1 > max_width - 1:
            lines.append(current_line)
            current_line = ""
    return "\\\n".join(lines)


def get_base_command(args, output_dir):

    
    args.base_cmd = re.sub(r"[\\\n]+", " ", args.base_cmd)

    
    args.base_cmd = re.sub("--output_dir\s+[^\s]+", "", args.base_cmd)
    args.base_cmd += f" --output_dir {output_dir}"

    
    args.base_cmd = re.sub("--overwrite_output_dir\s+", "", args.base_cmd)
    args.base_cmd += " --overwrite_output_dir"

    return [sys.executable] + shlex.split(args.base_cmd)


def process_run_single(id, cmd, variation, output_dir, target_metric_key, metric_keys, verbose):

    
    
    
    if 0:
        import random
        from time import sleep

        sleep(0)
        return dict(
            {k: random.uniform(0, 100) for k in metric_keys},
            **{target_metric_key: random.choice([nan, 10.31, 100.2, 55.6666, 222.22222222])},
        )

    result = subprocess.run(cmd, capture_output=True, text=True)

    if verbose:
        print("STDOUT", result.stdout)
        print("STDERR", result.stderr)

    
    prefix = variation.replace(" ", "-")
    with open(Path(output_dir) / f"log.{prefix}.stdout.txt", "w") as f:
        f.write(result.stdout)
    with open(Path(output_dir) / f"log.{prefix}.stderr.txt", "w") as f:
        f.write(result.stderr)

    if result.returncode != 0:
        if verbose:
            print("failed")
        return {target_metric_key: nan}

    with io.open(f"{output_dir}/all_results.json", "r", encoding="utf-8") as f:
        metrics = json.load(f)

    
    return {k: v for k, v in metrics.items() if k in metric_keys}


def process_run(
    id,
    cmd,
    variation_key,
    variation,
    longest_variation_len,
    target_metric_key,
    report_metric_keys,
    repeat_times,
    output_dir,
    verbose,
):
    results = []
    metrics = []
    preamble = f"{id}: {variation:<{longest_variation_len}}"
    outcome = f"{preamble}: "
    metric_keys = set(report_metric_keys + [target_metric_key])
    for i in tqdm(range(repeat_times), desc=preamble, leave=False):
        single_run_metrics = process_run_single(
            id, cmd, variation, output_dir, target_metric_key, metric_keys, verbose
        )
        result = single_run_metrics[target_metric_key]
        if not math.isnan(result):
            metrics.append(single_run_metrics)
            results.append(result)
            outcome += "✓"
        else:
            outcome += "✘"
    outcome = f"\33[2K\r{outcome}"
    if len(metrics) > 0:
        mean_metrics = {k: fmean([x[k] for x in metrics]) for k in metrics[0].keys()}
        mean_target = round(mean_metrics[target_metric_key], 2)
        results_str = f"{outcome} {mean_target}"
        if len(metrics) > 1:
            results_str += f" {tuple(round(x, 2) for x in results)}"
        print(results_str)
        mean_metrics[variation_key] = variation
        return mean_metrics
    else:
        print(outcome)
        return {variation_key: variation, target_metric_key: nan}


def get_versions():
    properties = torch.cuda.get_device_properties(torch.device("cuda"))
    return f


def process_results(results, target_metric_key, report_metric_keys, base_variation, output_dir):

    df = pd.DataFrame(results)
    variation_key = "variation"
    diff_key = "diff_%"

    sentinel_value = nan
    if base_variation is not None and len(df[df[variation_key] == base_variation]):
        
        sentinel_value = df.loc[df[variation_key] == base_variation][target_metric_key].item()
    if math.isnan(sentinel_value):
        
        sentinel_value = df.loc[df[target_metric_key] != nan][target_metric_key].min()

    
    if not math.isnan(sentinel_value):
        df[diff_key] = df.apply(
            lambda r: round(100 * (r[target_metric_key] - sentinel_value) / sentinel_value)
            if not math.isnan(r[target_metric_key])
            else 0,
            axis="columns",
        )

    
    cols = [variation_key, target_metric_key, diff_key, *report_metric_keys]
    df = df.reindex(cols, axis="columns")  

    
    df = df.rename(str.capitalize, axis="columns")

    
    df_github = df.rename(lambda c: c.replace("_", "<br>"), axis="columns")
    df_console = df.rename(lambda c: c.replace("_", "\n"), axis="columns")

    report = ["", "Copy between the cut-here-lines and paste as is to github or a forum"]
    report += ["----------8<-----------------8<--------"]
    report += ["*** Results:", df_github.to_markdown(index=False, floatfmt=".2f")]
    report += ["```"]
    report += ["*** Setup:", get_versions()]
    report += ["*** The benchmark command line was:", get_original_command()]
    report += ["```"]
    report += ["----------8<-----------------8<--------"]
    report += ["*** Results (console):", df_console.to_markdown(index=False, floatfmt=".2f")]

    print("\n\n".join(report))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--base-cmd",
        default=None,
        type=str,
        required=True,
        help="Base cmd",
    )
    parser.add_argument(
        "--variations",
        default=None,
        type=str,
        nargs="+",
        required=True,
        help="Multi-dimensional variations, example: '|--fp16|--bf16' '|--tf32'",
    )
    parser.add_argument(
        "--base-variation",
        default=None,
        type=str,
        help="Baseline variation to compare to. if None the minimal target value will be used to compare against",
    )
    parser.add_argument(
        "--target-metric-key",
        default=None,
        type=str,
        required=True,
        help="Target metric key in output_dir/all_results.json, e.g., train_samples_per_second",
    )
    parser.add_argument(
        "--report-metric-keys",
        default="",
        type=str,
        help="Report metric keys - other metric keys from output_dir/all_results.json to report, e.g., train_loss. Use a single argument e.g., 'train_loss train_samples",
    )
    parser.add_argument(
        "--repeat-times",
        default=1,
        type=int,
        help="How many times to re-run each variation - an average will be reported",
    )
    parser.add_argument(
        "--output_dir",
        default="output_benchmark",
        type=str,
        help="The output directory where all the benchmark reports will go to and additionally this directory will be used to override --output_dir in the script that is being benchmarked",
    )
    parser.add_argument(
        "--verbose",
        default=False,
        action="store_true",
        help="Whether to show the outputs of each run or just the benchmark progress",
    )
    args = parser.parse_args()

    output_dir = args.output_dir
    Path(output_dir).mkdir(exist_ok=True)
    base_cmd = get_base_command(args, output_dir)

    
    dims = [list(map(str.strip, re.split(r"\|", x))) for x in args.variations]
    
    
    variations = list(map(str.strip, map(" ".join, itertools.product(*dims))))
    longest_variation_len = max(len(x) for x in variations)

    
    report_metric_keys = args.report_metric_keys.split()

    
    report_fn = f"benchmark-report-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}.txt"
    print(f"\nNote: each run's output is also logged under {output_dir}/log.*.std*.txt")
    print(f"and this script's output is also piped into {report_fn}")

    sys.stdout = Tee(report_fn)

    print(f"\n*** Running {len(variations)} benchmarks:")
    print(f"Base command: {' '.join(base_cmd)}")

    variation_key = "variation"
    results = []
    for id, variation in enumerate(tqdm(variations, desc="Total completion: ", leave=False)):
        cmd = base_cmd + variation.split()
        results.append(
            process_run(
                id + 1,
                cmd,
                variation_key,
                variation,
                longest_variation_len,
                args.target_metric_key,
                report_metric_keys,
                args.repeat_times,
                output_dir,
                args.verbose,
            )
        )

    process_results(results, args.target_metric_key, report_metric_keys, args.base_variation, output_dir)


if __name__ == "__main__":
    main()