forked from phoenix-oss/llama-stack-mirror
fix: Correctly parse algorithm_config when launching NVIDIA customization job; fix internal request handler (#2025)
# What does this PR do? This addresses 2 bugs I ran into when launching a fine-tuning job with the NVIDIA Adapter: 1. Session handling in `_make_request` helper function returns an error. ``` INFO: 127.0.0.1:55831 - "POST /v1/post-training/supervised-fine-tune HTTP/1.1" 500 Internal Server Error 16:11:45.643 [END] /v1/post-training/supervised-fine-tune [StatusCode.OK] (270.44ms) 16:11:45.643 [ERROR] Error executing endpoint route='/v1/post-training/supervised-fine-tune' method='post' Traceback (most recent call last): File "/Users/jgulabrai/Projects/forks/llama-stack/llama_stack/distribution/server/server.py", line 201, in endpoint return await maybe_await(value) File "/Users/jgulabrai/Projects/forks/llama-stack/llama_stack/distribution/server/server.py", line 161, in maybe_await return await value File "/Users/jgulabrai/Projects/forks/llama-stack/llama_stack/providers/remote/post_training/nvidia/post_training.py", line 408, in supervised_fine_tune response = await self._make_request( File "/Users/jgulabrai/Projects/forks/llama-stack/llama_stack/providers/remote/post_training/nvidia/post_training.py", line 98, in _make_request async with self.session.request(method, url, params=params, json=json, **kwargs) as response: File "/Users/jgulabrai/Projects/forks/llama-stack/.venv/lib/python3.10/site-packages/aiohttp/client.py", line 1425, in __aenter__ self._resp: _RetType = await self._coro File "/Users/jgulabrai/Projects/forks/llama-stack/.venv/lib/python3.10/site-packages/aiohttp/client.py", line 579, in _request handle = tm.start() File "/Users/jgulabrai/Projects/forks/llama-stack/.venv/lib/python3.10/site-packages/aiohttp/helpers.py", line 587, in start return self._loop.call_at(when, self.__call__) File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/asyncio/base_events.py", line 724, in call_at self._check_closed() File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/asyncio/base_events.py", line 510, in _check_closed raise RuntimeError('Event loop is closed') RuntimeError: Event loop is closed ``` Note: This only occurred when initializing the client like so: ``` client = LlamaStackClient( base_url="http://0.0.0.0:8321" ) response = client.post_training.supervised_fine_tune(...) # Returns error ``` I didn't run into this issue when using the library client: ``` client = LlamaStackAsLibraryClient("nvidia") client.initialize() response = client.post_training.supervised_fine_tune(...) # Works fine ``` 2. The `algorithm_config` param in `supervised_fine_tune` is parsed as a `dict` when run from unit tests, but a Pydantic model when invoked using the Llama Stack client. So, the call fails outside of unit tests: ``` INFO: 127.0.0.1:54024 - "POST /v1/post-training/supervised-fine-tune HTTP/1.1" 500 Internal Server Error 21:14:02.315 [END] /v1/post-training/supervised-fine-tune [StatusCode.OK] (71.18ms) 21:14:02.314 [ERROR] Error executing endpoint route='/v1/post-training/supervised-fine-tune' method='post' Traceback (most recent call last): File "/Users/jgulabrai/Projects/forks/llama-stack/llama_stack/distribution/server/server.py", line 205, in endpoint return await maybe_await(value) File "/Users/jgulabrai/Projects/forks/llama-stack/llama_stack/distribution/server/server.py", line 164, in maybe_await return await value File "/Users/jgulabrai/Projects/forks/llama-stack/llama_stack/providers/remote/post_training/nvidia/post_training.py", line 407, in supervised_fine_tune "adapter_dim": algorithm_config.get("adapter_dim"), File "/Users/jgulabrai/Projects/forks/llama-stack/.venv/lib/python3.10/site-packages/pydantic/main.py", line 891, in __getattr__ raise AttributeError(f'{type(self).__name__!r} object has no attribute {item!r}') AttributeError: 'LoraFinetuningConfig' object has no attribute 'get' ``` The code assumes `algorithm_config` should be `dict`, so I just handle both cases. [//]: # (If resolving an issue, uncomment and update the line below) [//]: # (Closes #[issue-number]) ## Test Plan 1. I ran a local Llama Stack server with the necessary env vars: ``` lama stack run llama_stack/templates/nvidia/run.yaml --port 8321 --env ... ``` And invoked `supervised_fine_tune` to confirm neither of the errors above occur. ``` client = LlamaStackClient( base_url="http://0.0.0.0:8321" ) response = client.post_training.supervised_fine_tune(...) ``` 2. I confirmed the unit tests still pass: `./scripts/unit-tests.sh tests/unit/providers/nvidia/test_supervised_fine_tuning.py` [//]: # (## Documentation) --------- Co-authored-by: Jash Gulabrai <jgulabrai@nvidia.com>
This commit is contained in:
parent
b5d8e44e81
commit
8713d67ce3
3 changed files with 83 additions and 66 deletions
|
@ -67,13 +67,18 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
self.timeout = aiohttp.ClientTimeout(total=config.timeout)
|
self.timeout = aiohttp.ClientTimeout(total=config.timeout)
|
||||||
# TODO: filter by available models based on /config endpoint
|
# TODO: filter by available models based on /config endpoint
|
||||||
ModelRegistryHelper.__init__(self, model_entries=_MODEL_ENTRIES)
|
ModelRegistryHelper.__init__(self, model_entries=_MODEL_ENTRIES)
|
||||||
self.session = aiohttp.ClientSession(headers=self.headers, timeout=self.timeout)
|
self.session = None
|
||||||
self.customizer_url = config.customizer_url
|
|
||||||
|
|
||||||
|
self.customizer_url = config.customizer_url
|
||||||
if not self.customizer_url:
|
if not self.customizer_url:
|
||||||
warnings.warn("Customizer URL is not set, using default value: http://nemo.test", stacklevel=2)
|
warnings.warn("Customizer URL is not set, using default value: http://nemo.test", stacklevel=2)
|
||||||
self.customizer_url = "http://nemo.test"
|
self.customizer_url = "http://nemo.test"
|
||||||
|
|
||||||
|
async def _get_session(self) -> aiohttp.ClientSession:
|
||||||
|
if self.session is None or self.session.closed:
|
||||||
|
self.session = aiohttp.ClientSession(headers=self.headers, timeout=self.timeout)
|
||||||
|
return self.session
|
||||||
|
|
||||||
async def _make_request(
|
async def _make_request(
|
||||||
self,
|
self,
|
||||||
method: str,
|
method: str,
|
||||||
|
@ -94,8 +99,9 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
if json and "Content-Type" not in request_headers:
|
if json and "Content-Type" not in request_headers:
|
||||||
request_headers["Content-Type"] = "application/json"
|
request_headers["Content-Type"] = "application/json"
|
||||||
|
|
||||||
|
session = await self._get_session()
|
||||||
for _ in range(self.config.max_retries):
|
for _ in range(self.config.max_retries):
|
||||||
async with self.session.request(method, url, params=params, json=json, **kwargs) as response:
|
async with session.request(method, url, params=params, json=json, **kwargs) as response:
|
||||||
if response.status >= 400:
|
if response.status >= 400:
|
||||||
error_data = await response.json()
|
error_data = await response.json()
|
||||||
raise Exception(f"API request failed: {error_data}")
|
raise Exception(f"API request failed: {error_data}")
|
||||||
|
@ -122,8 +128,8 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
jobs = []
|
jobs = []
|
||||||
for job in response.get("data", []):
|
for job in response.get("data", []):
|
||||||
job_id = job.pop("id")
|
job_id = job.pop("id")
|
||||||
job_status = job.pop("status", "unknown").lower()
|
job_status = job.pop("status", "scheduled").lower()
|
||||||
mapped_status = STATUS_MAPPING.get(job_status, "unknown")
|
mapped_status = STATUS_MAPPING.get(job_status, "scheduled")
|
||||||
|
|
||||||
# Convert string timestamps to datetime objects
|
# Convert string timestamps to datetime objects
|
||||||
created_at = (
|
created_at = (
|
||||||
|
@ -177,7 +183,7 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
)
|
)
|
||||||
|
|
||||||
api_status = response.pop("status").lower()
|
api_status = response.pop("status").lower()
|
||||||
mapped_status = STATUS_MAPPING.get(api_status, "unknown")
|
mapped_status = STATUS_MAPPING.get(api_status, "scheduled")
|
||||||
|
|
||||||
return NvidiaPostTrainingJobStatusResponse(
|
return NvidiaPostTrainingJobStatusResponse(
|
||||||
status=JobStatus(mapped_status),
|
status=JobStatus(mapped_status),
|
||||||
|
@ -239,6 +245,7 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
|
|
||||||
Supported models:
|
Supported models:
|
||||||
- meta/llama-3.1-8b-instruct
|
- meta/llama-3.1-8b-instruct
|
||||||
|
- meta/llama-3.2-1b-instruct
|
||||||
|
|
||||||
Supported algorithm configs:
|
Supported algorithm configs:
|
||||||
- LoRA, SFT
|
- LoRA, SFT
|
||||||
|
@ -284,10 +291,6 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
|
|
||||||
- LoRA config:
|
- LoRA config:
|
||||||
## NeMo customizer specific LoRA parameters
|
## NeMo customizer specific LoRA parameters
|
||||||
- adapter_dim: int - Adapter dimension
|
|
||||||
Default: 8 (supports powers of 2)
|
|
||||||
- adapter_dropout: float - Adapter dropout
|
|
||||||
Default: None (0.0-1.0)
|
|
||||||
- alpha: int - Scaling factor for the LoRA update
|
- alpha: int - Scaling factor for the LoRA update
|
||||||
Default: 16
|
Default: 16
|
||||||
Note:
|
Note:
|
||||||
|
@ -297,7 +300,7 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
User is informed about unsupported parameters via warnings.
|
User is informed about unsupported parameters via warnings.
|
||||||
"""
|
"""
|
||||||
# Map model to nvidia model name
|
# Map model to nvidia model name
|
||||||
# ToDo: only supports llama-3.1-8b-instruct now, need to update this to support other models
|
# See `_MODEL_ENTRIES` for supported models
|
||||||
nvidia_model = self.get_provider_model_id(model)
|
nvidia_model = self.get_provider_model_id(model)
|
||||||
|
|
||||||
# Check for unsupported method parameters
|
# Check for unsupported method parameters
|
||||||
|
@ -330,7 +333,7 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
},
|
},
|
||||||
"data_config": {"dataset_id", "batch_size"},
|
"data_config": {"dataset_id", "batch_size"},
|
||||||
"optimizer_config": {"lr", "weight_decay"},
|
"optimizer_config": {"lr", "weight_decay"},
|
||||||
"lora_config": {"type", "adapter_dim", "adapter_dropout", "alpha"},
|
"lora_config": {"type", "alpha"},
|
||||||
}
|
}
|
||||||
|
|
||||||
# Validate all parameters at once
|
# Validate all parameters at once
|
||||||
|
@ -389,16 +392,10 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
|
||||||
|
|
||||||
# Handle LoRA-specific configuration
|
# Handle LoRA-specific configuration
|
||||||
if algorithm_config:
|
if algorithm_config:
|
||||||
if isinstance(algorithm_config, dict) and algorithm_config.get("type") == "LoRA":
|
if algorithm_config.type == "LoRA":
|
||||||
warn_unsupported_params(algorithm_config, supported_params["lora_config"], "LoRA config")
|
warn_unsupported_params(algorithm_config, supported_params["lora_config"], "LoRA config")
|
||||||
job_config["hyperparameters"]["lora"] = {
|
job_config["hyperparameters"]["lora"] = {
|
||||||
k: v
|
k: v for k, v in {"alpha": algorithm_config.alpha}.items() if v is not None
|
||||||
for k, v in {
|
|
||||||
"adapter_dim": algorithm_config.get("adapter_dim"),
|
|
||||||
"alpha": algorithm_config.get("alpha"),
|
|
||||||
"adapter_dropout": algorithm_config.get("adapter_dropout"),
|
|
||||||
}.items()
|
|
||||||
if v is not None
|
|
||||||
}
|
}
|
||||||
else:
|
else:
|
||||||
raise NotImplementedError(f"Unsupported algorithm config: {algorithm_config}")
|
raise NotImplementedError(f"Unsupported algorithm config: {algorithm_config}")
|
||||||
|
|
|
@ -10,14 +10,17 @@ import warnings
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from llama_stack_client.types.algorithm_config_param import LoraFinetuningConfig
|
|
||||||
from llama_stack_client.types.post_training_supervised_fine_tune_params import (
|
|
||||||
TrainingConfig,
|
|
||||||
TrainingConfigDataConfig,
|
|
||||||
TrainingConfigEfficiencyConfig,
|
|
||||||
TrainingConfigOptimizerConfig,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
from llama_stack.apis.post_training.post_training import (
|
||||||
|
DataConfig,
|
||||||
|
DatasetFormat,
|
||||||
|
EfficiencyConfig,
|
||||||
|
LoraFinetuningConfig,
|
||||||
|
OptimizerConfig,
|
||||||
|
OptimizerType,
|
||||||
|
TrainingConfig,
|
||||||
|
)
|
||||||
|
from llama_stack.distribution.library_client import convert_pydantic_to_json_value
|
||||||
from llama_stack.providers.remote.post_training.nvidia.post_training import (
|
from llama_stack.providers.remote.post_training.nvidia.post_training import (
|
||||||
NvidiaPostTrainingAdapter,
|
NvidiaPostTrainingAdapter,
|
||||||
NvidiaPostTrainingConfig,
|
NvidiaPostTrainingConfig,
|
||||||
|
@ -66,11 +69,8 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
|
|
||||||
def test_customizer_parameters_passed(self):
|
def test_customizer_parameters_passed(self):
|
||||||
"""Test scenario 1: When an optional parameter is passed and value is correctly set."""
|
"""Test scenario 1: When an optional parameter is passed and value is correctly set."""
|
||||||
custom_adapter_dim = 32 # Different from default of 8
|
|
||||||
algorithm_config = LoraFinetuningConfig(
|
algorithm_config = LoraFinetuningConfig(
|
||||||
type="LoRA",
|
type="LoRA",
|
||||||
adapter_dim=custom_adapter_dim,
|
|
||||||
adapter_dropout=0.2,
|
|
||||||
apply_lora_to_mlp=True,
|
apply_lora_to_mlp=True,
|
||||||
apply_lora_to_output=True,
|
apply_lora_to_output=True,
|
||||||
alpha=16,
|
alpha=16,
|
||||||
|
@ -78,8 +78,15 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
||||||
)
|
)
|
||||||
|
|
||||||
data_config = TrainingConfigDataConfig(dataset_id="test-dataset", batch_size=16)
|
data_config = DataConfig(
|
||||||
optimizer_config = TrainingConfigOptimizerConfig(lr=0.0002)
|
dataset_id="test-dataset", batch_size=16, shuffle=False, data_format=DatasetFormat.instruct
|
||||||
|
)
|
||||||
|
optimizer_config = OptimizerConfig(
|
||||||
|
optimizer_type=OptimizerType.adam,
|
||||||
|
lr=0.0002,
|
||||||
|
weight_decay=0.01,
|
||||||
|
num_warmup_steps=100,
|
||||||
|
)
|
||||||
training_config = TrainingConfig(
|
training_config = TrainingConfig(
|
||||||
n_epochs=3,
|
n_epochs=3,
|
||||||
data_config=data_config,
|
data_config=data_config,
|
||||||
|
@ -95,7 +102,7 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
model="meta-llama/Llama-3.1-8B-Instruct",
|
model="meta-llama/Llama-3.1-8B-Instruct",
|
||||||
checkpoint_dir="",
|
checkpoint_dir="",
|
||||||
algorithm_config=algorithm_config,
|
algorithm_config=algorithm_config,
|
||||||
training_config=training_config,
|
training_config=convert_pydantic_to_json_value(training_config),
|
||||||
logger_config={},
|
logger_config={},
|
||||||
hyperparam_search_config={},
|
hyperparam_search_config={},
|
||||||
)
|
)
|
||||||
|
@ -114,7 +121,7 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
self._assert_request_params(
|
self._assert_request_params(
|
||||||
{
|
{
|
||||||
"hyperparameters": {
|
"hyperparameters": {
|
||||||
"lora": {"adapter_dim": custom_adapter_dim, "adapter_dropout": 0.2, "alpha": 16},
|
"lora": {"alpha": 16},
|
||||||
"epochs": 3,
|
"epochs": 3,
|
||||||
"learning_rate": 0.0002,
|
"learning_rate": 0.0002,
|
||||||
"batch_size": 16,
|
"batch_size": 16,
|
||||||
|
@ -130,8 +137,6 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
|
|
||||||
algorithm_config = LoraFinetuningConfig(
|
algorithm_config = LoraFinetuningConfig(
|
||||||
type="LoRA",
|
type="LoRA",
|
||||||
adapter_dim=16,
|
|
||||||
adapter_dropout=0.1,
|
|
||||||
apply_lora_to_mlp=True,
|
apply_lora_to_mlp=True,
|
||||||
apply_lora_to_output=True,
|
apply_lora_to_output=True,
|
||||||
alpha=16,
|
alpha=16,
|
||||||
|
@ -139,12 +144,16 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
||||||
)
|
)
|
||||||
|
|
||||||
data_config = TrainingConfigDataConfig(
|
data_config = DataConfig(
|
||||||
dataset_id=required_dataset_id, # Required parameter
|
dataset_id=required_dataset_id, batch_size=8, shuffle=False, data_format=DatasetFormat.instruct
|
||||||
batch_size=8,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
optimizer_config = TrainingConfigOptimizerConfig(lr=0.0001)
|
optimizer_config = OptimizerConfig(
|
||||||
|
optimizer_type=OptimizerType.adam,
|
||||||
|
lr=0.0001,
|
||||||
|
weight_decay=0.01,
|
||||||
|
num_warmup_steps=100,
|
||||||
|
)
|
||||||
|
|
||||||
training_config = TrainingConfig(
|
training_config = TrainingConfig(
|
||||||
n_epochs=1,
|
n_epochs=1,
|
||||||
|
@ -161,7 +170,7 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
model=required_model, # Required parameter
|
model=required_model, # Required parameter
|
||||||
checkpoint_dir="",
|
checkpoint_dir="",
|
||||||
algorithm_config=algorithm_config,
|
algorithm_config=algorithm_config,
|
||||||
training_config=training_config,
|
training_config=convert_pydantic_to_json_value(training_config),
|
||||||
logger_config={},
|
logger_config={},
|
||||||
hyperparam_search_config={},
|
hyperparam_search_config={},
|
||||||
)
|
)
|
||||||
|
@ -186,24 +195,24 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
|
|
||||||
def test_unsupported_parameters_warning(self):
|
def test_unsupported_parameters_warning(self):
|
||||||
"""Test that warnings are raised for unsupported parameters."""
|
"""Test that warnings are raised for unsupported parameters."""
|
||||||
data_config = TrainingConfigDataConfig(
|
data_config = DataConfig(
|
||||||
dataset_id="test-dataset",
|
dataset_id="test-dataset",
|
||||||
batch_size=8,
|
batch_size=8,
|
||||||
# Unsupported parameters
|
# Unsupported parameters
|
||||||
shuffle=True,
|
shuffle=True,
|
||||||
data_format="instruct",
|
data_format=DatasetFormat.instruct,
|
||||||
validation_dataset_id="val-dataset",
|
validation_dataset_id="val-dataset",
|
||||||
)
|
)
|
||||||
|
|
||||||
optimizer_config = TrainingConfigOptimizerConfig(
|
optimizer_config = OptimizerConfig(
|
||||||
lr=0.0001,
|
lr=0.0001,
|
||||||
weight_decay=0.01,
|
weight_decay=0.01,
|
||||||
# Unsupported parameters
|
# Unsupported parameters
|
||||||
optimizer_type="adam",
|
optimizer_type=OptimizerType.adam,
|
||||||
num_warmup_steps=100,
|
num_warmup_steps=100,
|
||||||
)
|
)
|
||||||
|
|
||||||
efficiency_config = TrainingConfigEfficiencyConfig(
|
efficiency_config = EfficiencyConfig(
|
||||||
enable_activation_checkpointing=True # Unsupported parameter
|
enable_activation_checkpointing=True # Unsupported parameter
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@ -230,15 +239,13 @@ class TestNvidiaParameters(unittest.TestCase):
|
||||||
checkpoint_dir="test-dir", # Unsupported parameter
|
checkpoint_dir="test-dir", # Unsupported parameter
|
||||||
algorithm_config=LoraFinetuningConfig(
|
algorithm_config=LoraFinetuningConfig(
|
||||||
type="LoRA",
|
type="LoRA",
|
||||||
adapter_dim=16,
|
|
||||||
adapter_dropout=0.1,
|
|
||||||
apply_lora_to_mlp=True,
|
apply_lora_to_mlp=True,
|
||||||
apply_lora_to_output=True,
|
apply_lora_to_output=True,
|
||||||
alpha=16,
|
alpha=16,
|
||||||
rank=16,
|
rank=16,
|
||||||
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
||||||
),
|
),
|
||||||
training_config=training_config,
|
training_config=convert_pydantic_to_json_value(training_config),
|
||||||
logger_config={"test": "value"}, # Unsupported parameter
|
logger_config={"test": "value"}, # Unsupported parameter
|
||||||
hyperparam_search_config={"test": "value"}, # Unsupported parameter
|
hyperparam_search_config={"test": "value"}, # Unsupported parameter
|
||||||
)
|
)
|
||||||
|
|
|
@ -10,14 +10,18 @@ import warnings
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from llama_stack_client.types.algorithm_config_param import LoraFinetuningConfig, QatFinetuningConfig
|
|
||||||
from llama_stack_client.types.post_training_supervised_fine_tune_params import (
|
|
||||||
TrainingConfig,
|
|
||||||
TrainingConfigDataConfig,
|
|
||||||
TrainingConfigOptimizerConfig,
|
|
||||||
)
|
|
||||||
|
|
||||||
from llama_stack.apis.models import Model, ModelType
|
from llama_stack.apis.models import Model, ModelType
|
||||||
|
from llama_stack.apis.post_training.post_training import (
|
||||||
|
DataConfig,
|
||||||
|
DatasetFormat,
|
||||||
|
LoraFinetuningConfig,
|
||||||
|
OptimizerConfig,
|
||||||
|
OptimizerType,
|
||||||
|
QATFinetuningConfig,
|
||||||
|
TrainingConfig,
|
||||||
|
)
|
||||||
|
from llama_stack.distribution.library_client import convert_pydantic_to_json_value
|
||||||
from llama_stack.providers.remote.inference.nvidia.nvidia import NVIDIAConfig, NVIDIAInferenceAdapter
|
from llama_stack.providers.remote.inference.nvidia.nvidia import NVIDIAConfig, NVIDIAInferenceAdapter
|
||||||
from llama_stack.providers.remote.post_training.nvidia.post_training import (
|
from llama_stack.providers.remote.post_training.nvidia.post_training import (
|
||||||
ListNvidiaPostTrainingJobs,
|
ListNvidiaPostTrainingJobs,
|
||||||
|
@ -121,7 +125,7 @@ class TestNvidiaPostTraining(unittest.TestCase):
|
||||||
"batch_size": 16,
|
"batch_size": 16,
|
||||||
"epochs": 2,
|
"epochs": 2,
|
||||||
"learning_rate": 0.0001,
|
"learning_rate": 0.0001,
|
||||||
"lora": {"adapter_dim": 16, "adapter_dropout": 0.1},
|
"lora": {"alpha": 16},
|
||||||
},
|
},
|
||||||
"output_model": "default/job-1234",
|
"output_model": "default/job-1234",
|
||||||
"status": "created",
|
"status": "created",
|
||||||
|
@ -132,8 +136,6 @@ class TestNvidiaPostTraining(unittest.TestCase):
|
||||||
|
|
||||||
algorithm_config = LoraFinetuningConfig(
|
algorithm_config = LoraFinetuningConfig(
|
||||||
type="LoRA",
|
type="LoRA",
|
||||||
adapter_dim=16,
|
|
||||||
adapter_dropout=0.1,
|
|
||||||
apply_lora_to_mlp=True,
|
apply_lora_to_mlp=True,
|
||||||
apply_lora_to_output=True,
|
apply_lora_to_output=True,
|
||||||
alpha=16,
|
alpha=16,
|
||||||
|
@ -141,10 +143,15 @@ class TestNvidiaPostTraining(unittest.TestCase):
|
||||||
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
||||||
)
|
)
|
||||||
|
|
||||||
data_config = TrainingConfigDataConfig(dataset_id="sample-basic-test", batch_size=16)
|
data_config = DataConfig(
|
||||||
|
dataset_id="sample-basic-test", batch_size=16, shuffle=False, data_format=DatasetFormat.instruct
|
||||||
|
)
|
||||||
|
|
||||||
optimizer_config = TrainingConfigOptimizerConfig(
|
optimizer_config = OptimizerConfig(
|
||||||
|
optimizer_type=OptimizerType.adam,
|
||||||
lr=0.0001,
|
lr=0.0001,
|
||||||
|
weight_decay=0.01,
|
||||||
|
num_warmup_steps=100,
|
||||||
)
|
)
|
||||||
|
|
||||||
training_config = TrainingConfig(
|
training_config = TrainingConfig(
|
||||||
|
@ -161,7 +168,7 @@ class TestNvidiaPostTraining(unittest.TestCase):
|
||||||
model="meta-llama/Llama-3.1-8B-Instruct",
|
model="meta-llama/Llama-3.1-8B-Instruct",
|
||||||
checkpoint_dir="",
|
checkpoint_dir="",
|
||||||
algorithm_config=algorithm_config,
|
algorithm_config=algorithm_config,
|
||||||
training_config=training_config,
|
training_config=convert_pydantic_to_json_value(training_config),
|
||||||
logger_config={},
|
logger_config={},
|
||||||
hyperparam_search_config={},
|
hyperparam_search_config={},
|
||||||
)
|
)
|
||||||
|
@ -185,16 +192,22 @@ class TestNvidiaPostTraining(unittest.TestCase):
|
||||||
"epochs": 2,
|
"epochs": 2,
|
||||||
"batch_size": 16,
|
"batch_size": 16,
|
||||||
"learning_rate": 0.0001,
|
"learning_rate": 0.0001,
|
||||||
"lora": {"alpha": 16, "adapter_dim": 16, "adapter_dropout": 0.1},
|
"weight_decay": 0.01,
|
||||||
|
"lora": {"alpha": 16},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_supervised_fine_tune_with_qat(self):
|
def test_supervised_fine_tune_with_qat(self):
|
||||||
algorithm_config = QatFinetuningConfig(type="QAT", quantizer_name="quantizer_name", group_size=1)
|
algorithm_config = QATFinetuningConfig(type="QAT", quantizer_name="quantizer_name", group_size=1)
|
||||||
data_config = TrainingConfigDataConfig(dataset_id="sample-basic-test", batch_size=16)
|
data_config = DataConfig(
|
||||||
optimizer_config = TrainingConfigOptimizerConfig(
|
dataset_id="sample-basic-test", batch_size=16, shuffle=False, data_format=DatasetFormat.instruct
|
||||||
|
)
|
||||||
|
optimizer_config = OptimizerConfig(
|
||||||
|
optimizer_type=OptimizerType.adam,
|
||||||
lr=0.0001,
|
lr=0.0001,
|
||||||
|
weight_decay=0.01,
|
||||||
|
num_warmup_steps=100,
|
||||||
)
|
)
|
||||||
training_config = TrainingConfig(
|
training_config = TrainingConfig(
|
||||||
n_epochs=2,
|
n_epochs=2,
|
||||||
|
@ -209,7 +222,7 @@ class TestNvidiaPostTraining(unittest.TestCase):
|
||||||
model="meta-llama/Llama-3.1-8B-Instruct",
|
model="meta-llama/Llama-3.1-8B-Instruct",
|
||||||
checkpoint_dir="",
|
checkpoint_dir="",
|
||||||
algorithm_config=algorithm_config,
|
algorithm_config=algorithm_config,
|
||||||
training_config=training_config,
|
training_config=convert_pydantic_to_json_value(training_config),
|
||||||
logger_config={},
|
logger_config={},
|
||||||
hyperparam_search_config={},
|
hyperparam_search_config={},
|
||||||
)
|
)
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue