File size: 2,811 Bytes
a311e94 | 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 | """Data-order algorithms: sequence-length curriculum."""
from __future__ import annotations
import math
import random
from typing import List, Optional
import torch
from .base import AlgorithmContext, BaseAlgorithm
class Curriculum(BaseAlgorithm):
"""Sequence-length curriculum.
Early in training the engine is instructed to use short sequences; the
usable length then grows to the model maximum. This is a simple and
reliable curriculum for language modelling.
"""
name = "curriculum"
category = "data"
description = "Grows the training sequence length from short to full."
def __init__(self, start_ratio: float = 0.25, end_ratio: float = 1.0, schedule: str = "linear"):
super().__init__(start_ratio=start_ratio, end_ratio=end_ratio, schedule=schedule)
self.start_ratio = start_ratio
self.end_ratio = end_ratio
self.schedule = schedule
def _ratio(self, ctx) -> float:
total = max(1, ctx.max_steps)
t = min(1.0, ctx.step / total)
if self.schedule == "cosine":
t = 0.5 * (1.0 - math.cos(math.pi * t))
return self.start_ratio + (self.end_ratio - self.start_ratio) * t
def on_step_begin(self, ctx, batch):
ratio = self._ratio(ctx)
input_ids = batch.get("input_ids")
if input_ids is None or ratio >= 1.0:
ctx.set("seq_len_ratio", 1.0)
return batch
max_len = max(1, int(round(input_ids.size(1) * ratio)))
trimmed = dict(batch)
trimmed["input_ids"] = input_ids[:, :max_len]
if "labels" in batch and batch["labels"] is not None:
trimmed["labels"] = batch["labels"][:, :max_len]
if "attention_mask" in batch and batch["attention_mask"] is not None:
trimmed["attention_mask"] = batch["attention_mask"][:, :max_len]
ctx.set("seq_len_ratio", ratio)
return trimmed
class DifficultyCurriculum(BaseAlgorithm):
"""Sorts the dataset by a difficulty proxy (sequence length).
The engine exposes the active sampler through ``ctx.get("sampler")``.
"""
name = "difficulty_curriculum"
category = "data"
description = "Orders samples from short/easy to long/hard."
def __init__(self, chunk_size: int = 64, easy_first: bool = True):
super().__init__(chunk_size=chunk_size, easy_first=easy_first)
self.chunk_size = chunk_size
self.easy_first = easy_first
def on_train_begin(self, ctx):
dataset = ctx.get("dataset")
if dataset is None:
return
try:
lengths = [len(x) for x in dataset]
except Exception:
return
order = sorted(range(len(lengths)), key=lambda i: lengths[i], reverse=not self.easy_first)
ctx.set("sample_order", order)
|