Do not raise error when there is mixed dtype in the checkpoint (#6247)

Lunwen He · facebook-github-bot · commit 41ffa93d550d · 2024-10-15T18:10:15.000-07:00
Summary: Pull Request resolved: #6247 This warning message got turned into an exception during the refactor of D64145852. It fails our export flow for llama 3.2 1B/3B. This diff fixes it. Reviewed By: mergennachin, dvorjackz Differential Revision: D64438373 fbshipit-source-id: 3dbfe44b95c94111de8cb2773ee9e642f9b896e6
diff --git a/examples/models/checkpoint.py b/examples/models/checkpoint.py
@@ -67,7 +67,7 @@ def get_checkpoint_dtype(checkpoint: Dict[str, Any]) -> Optional[str]:
             if value.dtype != dtype
         ]
         if len(mismatched_dtypes) > 0:
-            raise ValueError(
+            print(
                 f"Mixed dtype model. Dtype of {first_key}: {first.dtype}. Mismatches in the checkpoint: {mismatched_dtypes}"
             )
     return dtype

Original file line number	Diff line number	Diff line change
`@@ -67,7 +67,7 @@ def get_checkpoint_dtype(checkpoint: Dict[str, Any]) -> Optional[str]:`
`67`	`67`	`if value.dtype != dtype`
`68`	`68`	`]`
`69`	`69`	`if len(mismatched_dtypes) > 0:`
`70`		`- raise ValueError(`
	`70`	`+ print(`
`71`	`71`	`f"Mixed dtype model. Dtype of {first_key}: {first.dtype}. Mismatches in the checkpoint: {mismatched_dtypes}"`
`72`	`72`	`)`
`73`	`73`	`return dtype`