Fix sharding config file name bug (#86)

* fix sharding config file name bug

* formmat

Author: FanhaiLu1

jetstream_pt/engine.py

@@ -711,7 +711,10 @@ def create_pytorch_engine(
   pt_model = None
 
   if not sharding_config:
-    sharding_config = os.path.join("default_shardings", model_name + ".yaml")
+    sharding_file_name = "llama" if model_name.startswith("llama") else "gemma"
+    sharding_config = os.path.join(
+        "default_shardings", sharding_file_name + ".yaml"
+    )
 
   env_data = JetEngineEnvironmentData(
       tokenizer_path=tokenizer_path,

run_server.py

@@ -105,10 +105,6 @@ def main(argv: Sequence[str]):
   devices = server_lib.get_devices()
   print(f"devices: {devices}")
   sharding_config_path = _SHARDING_CONFIG.value
-  if not sharding_config_path:
-    sharding_config_path = os.path.join(
-        "default_shardings", _MODEL_NAME.value + ".yaml"
-    )
   engine = jetstream_pt.create_pytorch_engine(
       devices=devices,
       tokenizer_path=_TOKENIZER_PATH.value,