File size: 2,147 Bytes
c9ac3c3 | 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 | # Copyright 2024 Bytedance Ltd. 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.
"""
Critic config
"""
from dataclasses import dataclass, field
from typing import Optional
from ..actor.config import FSDPConfig, ModelConfig, OffloadConfig, OptimConfig
@dataclass
class CriticConfig:
strategy: str = "fsdp"
global_batch_size: int = 256
"""number of samples per minibatch for updating critic"""
micro_batch_size_per_device_for_update: int = 4
"""number of samples per forward pass for updating critic"""
micro_batch_size_per_device_for_experience: int = 16
"""number of samples per forward pass for computing values"""
max_grad_norm: float = 1.0
"""number to clip grad norm"""
cliprange_value: float = 0.5
"""clip range for value loss"""
loss_avg_mode: str = "token"
"""loss average mode: `token`, `seq`"""
ppo_epochs: int = 1
"""number of ppo epochs for each rollout batch"""
padding_free: bool = False
"""use padding-free training"""
dynamic_batching: bool = True
"""enable dynamic batching"""
max_token_len_per_gpu: Optional[int] = None
"""max token length per GPU for dynamic batching. If None, use micro_batch_size * max_seq_len"""
ulysses_size: int = 1
"""ulysses sequence parallel size"""
model: ModelConfig = field(default_factory=ModelConfig)
optim: OptimConfig = field(default_factory=OptimConfig)
fsdp: FSDPConfig = field(default_factory=FSDPConfig)
offload: OffloadConfig = field(default_factory=OffloadConfig)
# below are auto keys
global_batch_size_per_device: int = field(default=-1, init=False)
|