Do not raise error when there is mixed dtype in the checkpoint

Lunwen He · facebook-github-bot · commit acb43d772dba · 2024-10-15T16:40:28.000-07:00
Summary: 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.

Differential Revision: D64438373
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`