cast_based_on_args() — pytorch Function Reference

Architecture documentation for the cast_based_on_args() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking calls 4 called by 3

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ cast_based_on_args() — pytorch Function Reference

Dependency Diagram

graph TD
  da8c81f0_c44b_4b89_914d_318c84a98a42["cast_based_on_args()"]
  684ace56_6d9f_5ea1_f47a_cb8aa5834a49["validate_model()"]
  684ace56_6d9f_5ea1_f47a_cb8aa5834a49 -->|calls| da8c81f0_c44b_4b89_914d_318c84a98a42
  da0c865a_ac14_7a10_8fc5_8a3b7509426d["maybe_cast()"]
  da0c865a_ac14_7a10_8fc5_8a3b7509426d -->|calls| da8c81f0_c44b_4b89_914d_318c84a98a42
  c9be2096_e6d7_2374_ad2e_a6e33f435ada["run()"]
  c9be2096_e6d7_2374_ad2e_a6e33f435ada -->|calls| da8c81f0_c44b_4b89_914d_318c84a98a42
  f77ea0ed_facd_a7d6_7e4a_17f4bae6c1fe["setup_amp()"]
  da8c81f0_c44b_4b89_914d_318c84a98a42 -->|calls| f77ea0ed_facd_a7d6_7e4a_17f4bae6c1fe
  966e8ab0_98c7_2f6a_55eb_0fa446f905f8["cast_to_fp32()"]
  da8c81f0_c44b_4b89_914d_318c84a98a42 -->|calls| 966e8ab0_98c7_2f6a_55eb_0fa446f905f8
  9e3d0bf0_1f54_f3db_4072_b52d57cefb52["cast_to_fp16()"]
  da8c81f0_c44b_4b89_914d_318c84a98a42 -->|calls| 9e3d0bf0_1f54_f3db_4072_b52d57cefb52
  e789c0d4_b1f7_d69d_736c_fb5c020174d2["cast_to_bf16()"]
  da8c81f0_c44b_4b89_914d_318c84a98a42 -->|calls| e789c0d4_b1f7_d69d_736c_fb5c020174d2
  style da8c81f0_c44b_4b89_914d_318c84a98a42 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 1993–2025

    def cast_based_on_args(self, model, example_inputs):
        if self.args.float32 or self.args.only in self.fp32_only_models:
            if not self.args.float32:
                log.warning("Model %s supports float32 only", self.args.only)
            model, example_inputs = cast_to_fp32(model, example_inputs)
        elif self.args.float16:
            if self.args.only in self.force_amp_for_fp16_bf16_models:
                log.warning(
                    "Model %s does not support float16, running with amp instead",
                    self.args.only,
                )
                self.args.amp = True
                self.setup_amp()
            else:
                model, example_inputs = cast_to_fp16(model, example_inputs)
        elif self.args.bfloat16:
            if self.args.only in self.force_amp_for_fp16_bf16_models:
                log.warning(
                    "Model %s does not support bfloat16, running with amp instead",
                    self.args.only,
                )
                self.args.amp = True
                self.setup_amp()
            elif self.args.only in self.force_fp16_for_bf16_models:
                log.warning(
                    "Model %s does not support bfloat16, running with float16 instead",
                    self.args.only,
                )
                model, example_inputs = cast_to_fp16(model, example_inputs)
            else:
                model, example_inputs = cast_to_bf16(model, example_inputs)

        return model, example_inputs