File size: 4,701 Bytes
4e046d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | # Copyright 2025 the LlamaFactory team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
from copy import deepcopy
from typing import Any
def _normalize_precision_enabled(value: Any) -> bool | str:
if isinstance(value, str):
value_lower = value.lower()
if value_lower == "true":
return True
if value_lower == "false":
return False
if value_lower == "auto":
return "auto"
return value
def infer_deepspeed_mixed_precision(ds_config: dict[str, Any]) -> str:
ds_config.setdefault("fp16", {})
ds_config.setdefault("bf16", {})
fp16_enabled = _normalize_precision_enabled(ds_config["fp16"].get("enabled", "auto"))
bf16_enabled = _normalize_precision_enabled(ds_config["bf16"].get("enabled", "auto"))
# This project only supports DeepSpeed bf16 or no mixed precision.
if fp16_enabled is True:
raise ValueError("DeepSpeed only supports bf16 mixed precision for now, fp16 is not supported.")
if bf16_enabled is True:
mixed_precision = "bf16"
elif bf16_enabled is False:
mixed_precision = "no"
elif fp16_enabled is False:
mixed_precision = "no"
else:
# When both bf16/fp16 are left as auto (or absent), default to bf16.
mixed_precision = "bf16"
ds_config["fp16"]["enabled"] = False
ds_config["bf16"]["enabled"] = mixed_precision == "bf16"
return mixed_precision
def _unset_hf_deepspeed_config() -> None:
try:
from transformers.integrations import unset_hf_deepspeed_config
except ImportError:
from transformers.deepspeed import unset_hf_deepspeed_config
unset_hf_deepspeed_config()
def _load_deepspeed_config(config_file: str) -> dict[str, Any]:
with open(config_file, encoding="utf-8") as f:
return json.load(f)
def setup_deepspeed_zero3_model_loading(is_train: bool, dist_config: dict[str, Any] | None):
"""Enable transformers' ZeRO-3-aware model loading for the current thread."""
config_file = dist_config.get("config_file")
if not config_file:
raise ValueError("DeepSpeed config_file is required in dist_config")
from accelerate.utils import DeepSpeedPlugin
try:
from transformers.integrations import is_deepspeed_zero3_enabled
except ImportError:
from transformers.deepspeed import is_deepspeed_zero3_enabled
# DeepSpeed configs often use "auto" placeholders that only make sense once
# we know the current runtime batch settings and precision mode.
ds_config = deepcopy(_load_deepspeed_config(config_file))
if "gradient_accumulation_steps" not in ds_config or ds_config["gradient_accumulation_steps"] == "auto":
ds_config["gradient_accumulation_steps"] = 1
if "train_micro_batch_size_per_gpu" not in ds_config or ds_config["train_micro_batch_size_per_gpu"] == "auto":
ds_config["train_micro_batch_size_per_gpu"] = 1
if ds_config.get("train_batch_size") == "auto":
ds_config.pop("train_batch_size")
zero_stage = ds_config.get("zero_optimization", {}).get("stage")
if zero_stage != 3:
return None
# ZeRO-3 model loading needs concrete fp16/bf16 flags, not "auto".
mixed_precision = infer_deepspeed_mixed_precision(ds_config)
plugin = DeepSpeedPlugin(hf_ds_config=ds_config, zero3_init_flag=True)
if not plugin.hf_ds_config.is_zero3():
return None
# Reuse the same precision inference rule as the training-time DeepSpeed path
# so both model-loading and engine setup stay aligned.
plugin.set_mixed_precision(mixed_precision)
plugin.set_deepspeed_weakref()
if not is_deepspeed_zero3_enabled():
raise RuntimeError(
"DeepSpeed ZeRO-3 model-loading bootstrap failed: transformers still reports zero3 disabled "
"after constructing HfDeepSpeedConfig. This usually means the runtime is using a different transformers "
"installation than expected, or the DeepSpeed global state was not established correctly."
)
return plugin
def teardown_deepspeed_zero3_model_loading(plugin) -> None:
if plugin is not None:
_unset_hf_deepspeed_config()
|