wehub-resource-sync
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
"""Abstract base for worker-side loss computation.
|
|
|
|
A ``Loss`` subclass defines ``forward_step`` and ``loss_func`` -- the two
|
|
methods that Megatron's pipeline-parallel scheduler calls during
|
|
training. Users who want to customise loss computation only need to
|
|
subclass ``Loss`` and override these two methods; no understanding of
|
|
the internal Megatron trainer is required.
|
|
|
|
Example::
|
|
|
|
class MyLoss(Loss):
|
|
def __init__(self, args):
|
|
self.label_smoothing = args.label_smoothing
|
|
|
|
def forward_step(self, data_iterator, model):
|
|
batch = next(data_iterator)
|
|
output = model(batch['input_ids'], ...)
|
|
return output, partial(self.loss_func, labels=batch['labels'])
|
|
|
|
def loss_func(self, output_tensor, *, labels):
|
|
loss = F.cross_entropy(output_tensor, labels)
|
|
return loss, {'loss': loss.item()}
|
|
|
|
Then register it::
|
|
|
|
register_ray_trainer('my_algo', trainer='...MyDriver', loss='...MyLoss')
|
|
"""
|
|
from abc import ABC, abstractmethod
|
|
|
|
|
|
class Loss(ABC):
|
|
"""Abstract base for worker-side loss / forward computation.
|
|
|
|
Mirrors the two methods that Megatron's PP scheduler calls:
|
|
``forward_step(data_iterator, model) -> (output, loss_fn)``
|
|
and ``loss_func(output_tensor, **ctx) -> (loss, metrics)``.
|
|
|
|
Subclasses may wrap an existing trainer via composition for code
|
|
reuse (see ``GRPOLoss``) or implement these from scratch.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def forward_step(self, data_iterator, model):
|
|
"""Run a single forward micro-batch through *model*.
|
|
|
|
Returns ``(output_tensor, partial(self.loss_func, ...))``.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def loss_func(self, output_tensor, **kwargs):
|
|
"""Compute scalar loss + metrics from ``output_tensor``.
|
|
|
|
Returns ``(loss, metric_dict)``.
|
|
"""
|