add tests.cal_flops.py

Former-commit-id: 47a119db6c6e937f6ed96f70e3cda6031b9fbd0d
2023-09-16 23:40:41 +08:00
parent 6443fef31a
commit 32f8b1082b
2 changed files with 46 additions and 4 deletions
--- a/src/llmtuner/tuner/core/loader.py
+++ b/src/llmtuner/tuner/core/loader.py
@@ -86,10 +86,8 @@ def load_model_and_tokenizer(

    # Fix config (for Qwen)
    if is_trainable and hasattr(config, "fp16") and hasattr(config, "bf16"):
-        if model_args.compute_dtype == torch.bfloat16:
-            setattr(config, "bf16", True)
-        else:
-            setattr(config, "fp16", True)
+        setattr(config, "fp16", model_args.compute_dtype == torch.float16)
+        setattr(config, "bf16", model_args.compute_dtype == torch.bfloat16)

    # Set RoPE scaling
    if model_args.rope_scaling is not None: