This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,236 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import json
|
||||
import os
|
||||
from typing import List, Literal, Optional, Tuple
|
||||
|
||||
from swift.template import ContextType, Messages, get_last_user_round
|
||||
from .utils import calculate_loss_scale
|
||||
|
||||
ALL_BASE_STRATEGY = ['default', 'last_round', 'all']
|
||||
|
||||
|
||||
class LossScale:
|
||||
"""Base class for loss scaling in training.
|
||||
|
||||
This class provides a flexible framework for controlling loss computation weights
|
||||
across different parts of the context (e.g., system prompts, user queries, assistant
|
||||
responses) during model training. Different strategies can be applied to selectively
|
||||
train on specific portions.
|
||||
|
||||
Attributes:
|
||||
is_binary (bool, optional): Indicates whether loss_scale contains only 0 and 1.
|
||||
If True, loss_scale will be replaced by labels to stay compatible with
|
||||
acceleration techniques such as liger_kernel.
|
||||
If False, an additional 'loss_scale' key will be stored and the
|
||||
corresponding loss function will be used.
|
||||
base_strategy (str): Base strategy for loss computation. One of 'default',
|
||||
'last_round', or 'all'.
|
||||
- 'default': Only compute loss on assistant responses
|
||||
- 'last_round': Only compute loss on the last round's assistant response
|
||||
- 'all': Compute loss on all parts
|
||||
"""
|
||||
is_binary = True
|
||||
|
||||
def __init__(self, base_strategy: Literal['default', 'last_round', 'all'] = 'default'):
|
||||
"""Initialize the loss scale object.
|
||||
|
||||
Args:
|
||||
base_strategy: Base strategy for loss computation. One of 'default',
|
||||
'last_round', or 'all'. Defaults to 'default'.
|
||||
|
||||
Raises:
|
||||
ValueError: If the provided base_strategy is not in the allowed list.
|
||||
"""
|
||||
if base_strategy not in ALL_BASE_STRATEGY:
|
||||
raise ValueError(f'ALL_BASE_STRATEGY: {ALL_BASE_STRATEGY}, base_strategy: {base_strategy}')
|
||||
self.base_strategy = base_strategy
|
||||
|
||||
def get_loss_scale(self, context: str, **kwargs) -> Tuple[List[str], List[float]]:
|
||||
"""Calculate loss scale for the given context.
|
||||
|
||||
This is a base implementation that subclasses can override to implement
|
||||
custom loss scaling logic.
|
||||
|
||||
Args:
|
||||
context: The input context (string).
|
||||
**kwargs: Additional keyword arguments, such as query (the query of the
|
||||
current round).
|
||||
|
||||
Returns:
|
||||
A tuple containing:
|
||||
- List[str]: List of contexts, potentially split into multiple parts
|
||||
- List[float]: Corresponding loss scale values, one-to-one with contexts
|
||||
"""
|
||||
return [context], [1.]
|
||||
|
||||
def __call__(self, context_list: List[str], context_types: List[ContextType], messages: Messages,
|
||||
**kwargs) -> Tuple[List[str], List[float]]:
|
||||
"""Process the complete conversation context and return contexts with loss scales.
|
||||
|
||||
This method iterates through all context segments and determines the loss scale
|
||||
for each based on the context type and base strategy. It handles special cases
|
||||
such as explicitly specified loss values in messages and pre-computed loss scales.
|
||||
|
||||
Args:
|
||||
context_list: List of context strings or dicts, each representing a segment
|
||||
of the conversation.
|
||||
context_types: List of context types corresponding to each context, indicating
|
||||
whether it's a system prompt, user query, assistant response, etc.
|
||||
messages: Complete message list containing the conversation history.
|
||||
|
||||
Returns:
|
||||
A tuple containing:
|
||||
- List[str]: Processed context list, potentially expanded if contexts
|
||||
are split into multiple parts
|
||||
- List[float]: Loss scale values corresponding one-to-one with the
|
||||
returned context list
|
||||
"""
|
||||
res_context_list = []
|
||||
res_loss_scale = []
|
||||
i = 0
|
||||
last_user_round = get_last_user_round(messages)
|
||||
for context, context_type in zip(context_list, context_types):
|
||||
is_last_round = 2 * i >= last_user_round
|
||||
query, loss, loss_scale = None, None, None
|
||||
if context_type == ContextType.RESPONSE:
|
||||
query = messages[2 * i]['content']
|
||||
# Currently, we only support applying loss/mask to the response part.
|
||||
loss = messages[2 * i + 1].get('loss')
|
||||
loss_scale = messages[2 * i + 1].get('loss_scale')
|
||||
assert context == messages[2 * i + 1]['content']
|
||||
i += 1
|
||||
if not isinstance(context, list) or (len(context) > 0 and isinstance(context[0], int)):
|
||||
context = [context]
|
||||
for j in range(len(context)):
|
||||
new_context, new_loss_scale = self._inner_call(
|
||||
context[j], context_type, query, loss[j] if isinstance(loss, list) else loss,
|
||||
loss_scale[j] if isinstance(loss_scale, list) else loss_scale, is_last_round)
|
||||
res_context_list += new_context
|
||||
res_loss_scale += new_loss_scale
|
||||
|
||||
# The values in loss_scale_list correspond one-to-one with the values in context_list.
|
||||
return res_context_list, res_loss_scale
|
||||
|
||||
def _inner_call(self, context, context_type, query, loss, loss_scale, is_last_round):
|
||||
if isinstance(context, dict) and 'loss_scale' in context:
|
||||
new_context = [[token] for token in context['token_ids']]
|
||||
loss_scale = context['loss_scale']
|
||||
else:
|
||||
if isinstance(context, dict) and 'token_ids' in context:
|
||||
context = context['token_ids']
|
||||
is_assistant = context_type in {ContextType.RESPONSE, ContextType.SUFFIX}
|
||||
if loss or loss is None and (self.base_strategy == 'all' or
|
||||
(self.base_strategy == 'default' and is_assistant) or
|
||||
(self.base_strategy == 'last_round' and is_assistant and is_last_round)):
|
||||
if loss_scale is None:
|
||||
new_context, loss_scale = self.get_loss_scale(context, query=query)
|
||||
else:
|
||||
new_context, loss_scale = [context], [loss_scale]
|
||||
else:
|
||||
new_context, loss_scale = [context], [0.]
|
||||
return new_context, loss_scale
|
||||
|
||||
@property
|
||||
def is_binary_loss_scale(self):
|
||||
"""Check if loss scale values are binary (only 0 and 1)."""
|
||||
return self.is_binary
|
||||
|
||||
|
||||
class ConfigLossScale(LossScale):
|
||||
"""Loss scale class that loads configuration from a JSON file.
|
||||
|
||||
This class extends LossScale to support loading predefined loss scale mappings
|
||||
from a configuration file. The mappings can specify different loss weights for
|
||||
different tokens or segments based on the content.
|
||||
|
||||
Attributes:
|
||||
loss_scale_config (str, optional): Path to the loss scale configuration file
|
||||
relative to the 'config' directory.
|
||||
loss_scale_map (dict, optional): Dictionary mapping loaded from the config file,
|
||||
containing predefined loss scale values for specific patterns or tokens.
|
||||
"""
|
||||
is_binary = None
|
||||
loss_scale_config = None # path
|
||||
|
||||
def __init__(self, base_strategy: Literal['default', 'last_round', 'all'] = 'default'):
|
||||
"""Initialize the config-based loss scale object.
|
||||
|
||||
Loads the loss scale configuration from a JSON file if loss_scale_config
|
||||
is specified.
|
||||
|
||||
Args:
|
||||
base_strategy: Base strategy for loss computation. One of 'default',
|
||||
'last_round', or 'all'. Defaults to 'default'.
|
||||
"""
|
||||
super().__init__(base_strategy)
|
||||
self.loss_scale_map = None
|
||||
if self.loss_scale_config is not None:
|
||||
path = os.path.dirname(os.path.abspath(__file__))
|
||||
config_path = os.path.join(path, 'config', self.loss_scale_config)
|
||||
with open(config_path, 'r', encoding='utf-8') as json_file:
|
||||
self.loss_scale_map = json.load(json_file)
|
||||
|
||||
@property
|
||||
def is_binary_loss_scale(self):
|
||||
if self.is_binary is not None:
|
||||
return self.is_binary
|
||||
if self.loss_scale_map is None:
|
||||
return True
|
||||
return all(scale in {0.0, 1.0} for lst in self.loss_scale_map.values() for scale in lst)
|
||||
|
||||
def get_loss_scale(self, context: str, *, query: Optional[str] = None, **kwargs):
|
||||
"""Calculate loss scale using the loaded configuration.
|
||||
|
||||
If context is a string, uses the configuration map to calculate loss scales
|
||||
based on the query and context. Otherwise, falls back to the parent class
|
||||
implementation.
|
||||
|
||||
Args:
|
||||
context: The input context string.
|
||||
query: The user query for the current round, used to determine
|
||||
appropriate loss scaling based on the configuration.
|
||||
|
||||
Returns:
|
||||
Tuple[List[str], List[float]]: List of context segments and their
|
||||
corresponding loss scale values.
|
||||
"""
|
||||
if isinstance(context, str):
|
||||
return calculate_loss_scale(query, context, self.loss_scale_map)
|
||||
return super().get_loss_scale(context)
|
||||
|
||||
|
||||
class ConcatLossScale(LossScale):
|
||||
"""Apply multiple loss scales sequentially.
|
||||
|
||||
The output segments of each underlying loss scale are fed into the next one,
|
||||
and the corresponding weights are multiplied together. This makes it possible
|
||||
to compose strategies such as ``hermes+ignore_empty_think``.
|
||||
"""
|
||||
|
||||
is_binary = None
|
||||
|
||||
def __init__(self,
|
||||
loss_scales: List[LossScale],
|
||||
base_strategy: Literal['default', 'last_round', 'all'] = 'default'):
|
||||
super().__init__(base_strategy)
|
||||
assert loss_scales, 'loss_scales must be a non-empty list'
|
||||
self.loss_scales = loss_scales
|
||||
|
||||
def get_loss_scale(self, context, **kwargs):
|
||||
contexts = [context]
|
||||
weights = [1.0]
|
||||
for ls in self.loss_scales:
|
||||
new_contexts: List = []
|
||||
new_weights: List[float] = []
|
||||
for c, w in zip(contexts, weights):
|
||||
sub_contexts, sub_weights = ls.get_loss_scale(c, **kwargs)
|
||||
for sc, sw in zip(sub_contexts, sub_weights):
|
||||
new_contexts.append(sc)
|
||||
new_weights.append(w * sw)
|
||||
contexts = new_contexts
|
||||
weights = new_weights
|
||||
return contexts, weights
|
||||
|
||||
@property
|
||||
def is_binary_loss_scale(self):
|
||||
return all(ls.is_binary_loss_scale for ls in self.loss_scales)
|
||||
Reference in New Issue
Block a user