From 5321af81da155c0cb7a5428a29438c33ca679865 Mon Sep 17 00:00:00 2001 From: Aleksandr Borzunov Date: Wed, 24 Jul 2024 12:43:17 +0000 Subject: [PATCH] Remove excess load_state_dict() leading to meta tensor warnings --- src/petals/server/from_pretrained.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/petals/server/from_pretrained.py b/src/petals/server/from_pretrained.py index 4a3b15077..ac6a23e0c 100644 --- a/src/petals/server/from_pretrained.py +++ b/src/petals/server/from_pretrained.py @@ -64,10 +64,6 @@ def load_pretrained_block( max_disk_space=max_disk_space, ) - # dummy load, check that keys match - report = block.load_state_dict(state_dict, strict=False) - assert not report.missing_keys, f"Some block weights are missing: {report.missing_keys}" - for param_name, _ in block.named_parameters(): assert param_name in state_dict, f"{param_name} not in state dict" param = state_dict[param_name] @@ -76,7 +72,6 @@ def load_pretrained_block( set_module_tensor_to_device(block, param_name, "cpu", value=param, dtype=param.dtype) logger.info(f"Loaded {model_name} block {block_index}") - logger.debug(f"Details: {report}") return block