File size: 5,764 Bytes
4d20b62 | 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 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | # Copyright 2024 ByteDance and/or its affiliates.
#
# 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 inspect
import torch
import torch.distributed as dist
def get_adamw(
model: torch.nn.Module,
weight_decay: float,
learning_rate: float,
betas: tuple[float, float],
device_type: str,
) -> torch.optim.AdamW:
"""
Create an AdamW optimizer for the given model with specified parameters.
Args:
model (torch.nn.Module): The model for which the optimizer is created.
weight_decay (float): The weight decay (L2 penalty) for the optimizer.
learning_rate (float): The learning rate for the optimizer.
betas (tuple): Coefficients used for computing running averages of gradient and its square.
device_type (str): The device type ('cuda' or 'cpu') on which the optimizer will operate.
Returns:
torch.optim.AdamW: The AdamW optimizer configured with the specified parameters.
"""
# start with all of the candidate parameters
param_dict = {pn: p for pn, p in model.named_parameters()}
# filter out those that do not require grad
param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
# create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
# i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
optim_groups = [
{"params": decay_params, "weight_decay": weight_decay},
{"params": nodecay_params, "weight_decay": 0.0},
]
num_decay_params = sum(p.numel() for p in decay_params)
num_nodecay_params = sum(p.numel() for p in nodecay_params)
print(
f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters"
)
print(
f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters"
)
# Create AdamW optimizer and use the fused version if it is available
fused_available = "fused" in inspect.signature(torch.optim.AdamW).parameters
use_fused = fused_available and device_type == "cuda"
extra_args = dict(fused=True) if use_fused else dict()
optimizer = torch.optim.AdamW(
optim_groups, lr=learning_rate, betas=betas, **extra_args
)
print(f"using fused AdamW: {use_fused}")
return optimizer
def get_optimizer(
configs, model: torch.nn.Module, param_names=None
) -> torch.optim.Optimizer:
print("param_names: ", param_names)
if len(param_names) == 0 or param_names[0] == "":
param_names = None
if configs.adam.use_adamw:
optimizer = get_adamw(
model=model,
weight_decay=configs.adam.weight_decay,
learning_rate=configs.adam.lr,
other_learning_rate=configs.other_lr,
betas=(configs.adam.beta1, configs.adam.beta2),
device_type="cuda" if torch.cuda.is_available() else "cpu",
)
else:
if param_names is None or len(param_names) == 0:
param_groups = [{"params": model.parameters(), "lr": configs.lr}]
else:
finetune_params = []
finetune_params_names = []
other_params = []
for n, p in model.named_parameters():
if any([name in n for name in param_names]):
finetune_params.append(p)
finetune_params_names.append(n)
else:
other_params.append(p)
param_groups = [{"params": finetune_params, "lr": configs.finetune.lr}]
print("Sample Finetune Params Names: ", finetune_params_names[:5])
if configs.lr > 0.0:
param_groups.append({"params": other_params, "lr": configs.lr})
optimizer = torch.optim.Adam(
param_groups,
lr=configs.adam.lr,
weight_decay=configs.adam.weight_decay,
betas=(configs.adam.beta1, configs.adam.beta2),
)
for i, param_group in enumerate(optimizer.param_groups):
print(
f"Parameter Group {i}: Learning Rate = {param_group['lr']} Num = {len(param_group['params'])}"
)
return optimizer
def is_loss_nan_check(loss: torch.Tensor) -> bool:
"""check the validness of the current loss
Args:
loss: the loss from the model
Returns:
bool: if True, loss is not nan or inf
"""
def is_nan(x):
return torch.isnan(x).any() or torch.isinf(x).any()
def all_reduce_tensor(tensor, op=dist.ReduceOp.SUM):
if dist.is_initialized():
dist.all_reduce(tensor, op=op)
return tensor
nan_flag = torch.tensor(
1.0 if is_nan(loss) else 0.0,
device=loss.device if torch.cuda.is_available() else None,
) # support cpu
# avoid "Watchdog caught collective operation timeout" error
all_reduce_tensor(nan_flag)
if nan_flag.item() > 0.0:
return True
return False
|