wehub-resource-sync
237 lines
10 KiB
Python
237 lines
10 KiB
Python
# 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)
|