phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-03 01:03:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

address commit

Browse source
This commit is contained in:
Botao Chen 2024-12-13 10:38:53 -08:00
parent 0f78a5fb2d
commit 29d0896ec8
1 changed files with 8 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

29
llama_stack/providers/inline/post_training/torchtune/utils.py
View file

@ -38,27 +38,22 @@ class ModelConfig(BaseModel):
checkpoint_type: str checkpoint_type: str
class ModelConfigs(BaseModel):
Llama3_2_3B_Instruct: ModelConfig
Llama_3_8B_Instruct: ModelConfig
class DatasetSchema(BaseModel): class DatasetSchema(BaseModel):
alpaca: List[Dict[str, ParamType]] alpaca: List[Dict[str, ParamType]]
MODEL_CONFIGS = ModelConfigs( MODEL_CONFIGS: Dict[str, ModelConfig] = {
Llama3_2_3B_Instruct=ModelConfig( "Llama3.2-3B-Instruct": ModelConfig(
model_definition=lora_llama3_2_3b, model_definition=lora_llama3_2_3b,
tokenizer_type=llama3_tokenizer, tokenizer_type=llama3_tokenizer,
checkpoint_type="LLAMA3_2", checkpoint_type="LLAMA3_2",
), ),
Llama_3_8B_Instruct=ModelConfig( "Llama-3-8B-Instruct": ModelConfig(
model_definition=lora_llama3_8b, model_definition=lora_llama3_8b,
tokenizer_type=llama3_tokenizer, tokenizer_type=llama3_tokenizer,
checkpoint_type="LLAMA3", checkpoint_type="LLAMA3",
), ),
) }
EXPECTED_DATASET_SCHEMA = DatasetSchema( EXPECTED_DATASET_SCHEMA = DatasetSchema(
@ -85,14 +80,9 @@ BuildLoraModelCallable = Callable[..., torch.nn.Module]
BuildTokenizerCallable = Callable[..., Llama3Tokenizer] BuildTokenizerCallable = Callable[..., Llama3Tokenizer]
def _modify_model_id(model_id: str) -> str:
return model_id.replace("-", "_").replace(".", "_")
def _validate_model_id(model_id: str) -> Model: def _validate_model_id(model_id: str) -> Model:
model = resolve_model(model_id) model = resolve_model(model_id)
modified_model_id = _modify_model_id(model.core_model_id.value) if model is None or model.core_model_id.value not in MODEL_CONFIGS:
if model is None or not hasattr(MODEL_CONFIGS, modified_model_id):
raise ValueError(f"Model {model_id} is not supported.") raise ValueError(f"Model {model_id} is not supported.")
return model return model
@ -101,8 +91,7 @@ async def get_model_definition(
model_id: str, model_id: str,
) -> BuildLoraModelCallable: ) -> BuildLoraModelCallable:
model = _validate_model_id(model_id) model = _validate_model_id(model_id)
modified_model_id = _modify_model_id(model.core_model_id.value) model_config = MODEL_CONFIGS[model.core_model_id.value]
model_config = getattr(MODEL_CONFIGS, modified_model_id)
if not hasattr(model_config, "model_definition"): if not hasattr(model_config, "model_definition"):
raise ValueError(f"Model {model_id} does not have model definition.") raise ValueError(f"Model {model_id} does not have model definition.")
return model_config.model_definition return model_config.model_definition
@ -112,8 +101,7 @@ async def get_tokenizer_type(
model_id: str, model_id: str,
) -> BuildTokenizerCallable: ) -> BuildTokenizerCallable:
model = _validate_model_id(model_id) model = _validate_model_id(model_id)
modified_model_id = _modify_model_id(model.core_model_id.value) model_config = MODEL_CONFIGS[model.core_model_id.value]
model_config = getattr(MODEL_CONFIGS, modified_model_id)
if not hasattr(model_config, "tokenizer_type"): if not hasattr(model_config, "tokenizer_type"):
raise ValueError(f"Model {model_id} does not have tokenizer_type.") raise ValueError(f"Model {model_id} does not have tokenizer_type.")
return model_config.tokenizer_type return model_config.tokenizer_type
@ -127,8 +115,7 @@ async def get_checkpointer_model_type(
For example, llama3.2 model tied weights (https://github.com/pytorch/torchtune/blob/main/torchtune/training/checkpointing/_checkpointer.py#L1041) For example, llama3.2 model tied weights (https://github.com/pytorch/torchtune/blob/main/torchtune/training/checkpointing/_checkpointer.py#L1041)
""" """
model = _validate_model_id(model_id) model = _validate_model_id(model_id)
modified_model_id = _modify_model_id(model.core_model_id.value) model_config = MODEL_CONFIGS[model.core_model_id.value]
model_config = getattr(MODEL_CONFIGS, modified_model_id)
if not hasattr(model_config, "checkpoint_type"): if not hasattr(model_config, "checkpoint_type"):
raise ValueError(f"Model {model_id} does not have checkpoint_type.") raise ValueError(f"Model {model_id} does not have checkpoint_type.")
return model_config.checkpoint_type return model_config.checkpoint_type