Update on "Remove sharded ckpt from export_llama"

Sharded checkpoint isn't used anymore; removing it and simplifying export_llama.

Differential Revision: [D87828518](https://our.internmc.facebook.com/intern/diff/D87828518/)

[ghstack-poisoned]

Author: lucylq

examples/models/llama/export_llama_lib.py

@@ -672,18 +672,12 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager:
         if llm_config.base.checkpoint
         else None
     )
-    checkpoint_dir = (
-        canonical_path(llm_config.base.checkpoint_dir)
-        if llm_config.base.checkpoint_dir
-        else None
-    )
     params_path = (
         canonical_path(llm_config.base.params) if llm_config.base.params else None
     )
     output_dir_path = canonical_path(llm_config.export.output_dir, dir=True)
 
     llm_config.base.checkpoint = checkpoint_path
-    llm_config.base.checkpoint_dir = checkpoint_dir
     llm_config.base.params = params_path
     llm_config.export.output_dir = output_dir_path
 

examples/models/llama/model.py

@@ -67,6 +67,7 @@ def __init__(self, llm_config: Optional[LlmConfig] = None):
         # Follow the instruction in https://github.com/facebookresearch/llama to download the model.
         device = "cpu"
         # flake8: noqa: TOR102
+        checkpoint = {}
         if checkpoint_path:
             checkpoint = torch.load(checkpoint_path, map_location=device, mmap=True)