Validate the bf16 precisions across models

[xuzhao9]

A couple of models only run on the default precisions, and we would like to to enable bf16 precisions on them.

The reason models are failing because their inputs either can't be directly cast to bf16, or not completely being cast to bf16.

Per @drisspg:

Using this script:

from transformer_nuggets.utils.shape_trace import ShapeLog
import torch
from pathlib import Path
from tqdm import tqdm
import logging
import json

logging.basicConfig(level=logging.INFO)

def main():
    import torchbenchmark.models as models
    models = []
    success_count = 0
    failure_count = 0
    model_failures = {}
    for file in Path("torchbenchmark/models/").iterdir():
        if file.is_dir():
            models.append(file.name)
    for model_name in tqdm(models, desc="Logging models", aunit="model"):
        try:
            module = __import__(f"torchbenchmark.models.{model_name}", fromlist=[model_name])
            model, example_inputs = module.Model(test="train", device="cuda", extra_args=["--precision=bf16",]).get_module()
            model(*example_inputs)
            success_count += 1
        except Exception as e:
            tqdm.write(f"Failed to log {module}: {e}")
            failure_count += 1
            model_failures[model_name] = str(e)

    tqdm.write(f"Successfully logged {success_count} models")
    tqdm.write(f"Failed to log {failure_count} models")
    with open("model_failures_bf16.txt", "w") as f:
        json.dump(model_failures, f)



if __name__ == "__main__":
    main()

returns the following model failures;

internalfb.com/intern/paste/P1197341526

[xuzhao9]

We also want to validate amp_bf16 on CPU.

[xuzhao9]

We are migrating to pt2 benchmark runner, so we do not plan to support bf16 in the legacy runner.