tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

update tp validation

Browse Source
This commit is contained in:
Wing Lian
2025-07-22 21:20:57 -04:00
parent 8fe4758e94
commit 9a2da4d9f0
1 changed files with 23 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/axolotl/utils/schemas/validation.py
View File

@@ -913,15 +913,14 @@ class OptimizationValidationMixin:
def check_tensor_parallel_size_update_ds_json(cls, data): def check_tensor_parallel_size_update_ds_json(cls, data):
tensor_parallel_size = data.get("tensor_parallel_size") tensor_parallel_size = data.get("tensor_parallel_size")
if tensor_parallel_size is not None and tensor_parallel_size > 1: if tensor_parallel_size is not None and tensor_parallel_size > 1:
if not data.get("deepspeed"): if data.get("deepspeed"):
raise ValueError(
"Tensor parallelism (TP) is only supported with DeepSpeed"
)
with open(data.get("deepspeed"), "r", encoding="utf-8") as ds_fin: with open(data.get("deepspeed"), "r", encoding="utf-8") as ds_fin:
ds_config = json.load(ds_fin) ds_config = json.load(ds_fin)
should_save = False should_save = False
if "tensor_parallel" not in ds_config: if "tensor_parallel" not in ds_config:
ds_config["tensor_parallel"] = {"autotp_size": tensor_parallel_size} ds_config["tensor_parallel"] = {
"autotp_size": tensor_parallel_size
}
should_save = True should_save = True
if ( if (
"gather_16bit_weights_on_model_save" "gather_16bit_weights_on_model_save"