tracegrpo / src /rl /scale_sampler.py

Archive complete TraceGRPO project results and documentation (2026-10-09)

646ccfc verified about 1 month ago

2.58 kB

"""Epoch-wise full TRAIN coverage in 64-prompt, eight-generation groups."""
from __future__ import annotations
import json,math
from pathlib import Path
import torch
from torch.utils.data import Sampler


class ScaleGroupSampler(Sampler[int]):
    def __init__(self, size:int, *, seed:int=42, start_epoch:int=0,
                 prompts_per_step:int=64, generations:int=8,
                 dataloader_repeats:int=512, log_path:Path|None=None):
        if size<64 or prompts_per_step<=0 or generations<=0 or start_epoch<0 or dataloader_repeats<=0:
            raise ValueError('invalid scale sampler configuration')
        self.size=size;self.seed=seed;self.epoch=start_epoch
        self.prompts_per_step=prompts_per_step;self.generations=generations
        self.dataloader_repeats=dataloader_repeats
        self.log_path=log_path

    @property
    def steps_per_epoch(self):
        return math.ceil(self.size/self.prompts_per_step)

    def __len__(self):
        return self.steps_per_epoch*self.prompts_per_step*self.generations*self.dataloader_repeats

    def __iter__(self):
        epoch=self.epoch;self.epoch+=1
        generator=torch.Generator().manual_seed(self.seed+epoch)
        order=torch.randperm(self.size,generator=generator).tolist()
        extra=(-self.size)%self.prompts_per_step
        # Select padding from the epoch head. No prompt is repeated within
        # one64-prompt step, even for the short final group.
        if extra:order.extend(order[:extra])
        if self.log_path:
            self.log_path.parent.mkdir(parents=True,exist_ok=True)
            with self.log_path.open('a') as out:
                out.write(json.dumps(dict(epoch=epoch,shuffle_seed=self.seed+epoch,
                    train_rows=self.size,padding_reused_rows=extra,
                    steps=self.steps_per_epoch,
                    dataloader_repeats_per_optimizer_step=self.dataloader_repeats))+'\n')
        # TRL 0.21 consumes one 512-row dataloader batch for generation and
        # then 511 more dataloader batches while training the 512 buffered
        # microbatches. Those subsequent inputs are ignored. Repeat each
        # complete generation group so the next optimizer step advances to
        # the next group, rather than skipping 511 groups or ending an epoch.
        for start in range(0,len(order),self.prompts_per_step):
            group=order[start:start+self.prompts_per_step]
            for _ in range(self.dataloader_repeats):
                for index in group:
                    for _ in range(self.generations):yield index

	"""Epoch-wise full TRAIN coverage in 64-prompt, eight-generation groups."""
	from __future__ import annotations
	import json,math
	from pathlib import Path
	import torch
	from torch.utils.data import Sampler


	class ScaleGroupSampler(Sampler[int]):
	def __init__(self, size:int, *, seed:int=42, start_epoch:int=0,
	prompts_per_step:int=64, generations:int=8,
	dataloader_repeats:int=512, log_path:Path\|None=None):
	if size<64 or prompts_per_step<=0 or generations<=0 or start_epoch<0 or dataloader_repeats<=0:
	raise ValueError('invalid scale sampler configuration')
	self.size=size;self.seed=seed;self.epoch=start_epoch
	self.prompts_per_step=prompts_per_step;self.generations=generations
	self.dataloader_repeats=dataloader_repeats
	self.log_path=log_path

	@property
	def steps_per_epoch(self):
	return math.ceil(self.size/self.prompts_per_step)

	def __len__(self):
	return self.steps_per_epochself.prompts_per_stepself.generations*self.dataloader_repeats

	def __iter__(self):
	epoch=self.epoch;self.epoch+=1
	generator=torch.Generator().manual_seed(self.seed+epoch)
	order=torch.randperm(self.size,generator=generator).tolist()
	extra=(-self.size)%self.prompts_per_step
	# Select padding from the epoch head. No prompt is repeated within
	# one64-prompt step, even for the short final group.
	if extra:order.extend(order[:extra])
	if self.log_path:
	self.log_path.parent.mkdir(parents=True,exist_ok=True)
	with self.log_path.open('a') as out:
	out.write(json.dumps(dict(epoch=epoch,shuffle_seed=self.seed+epoch,
	train_rows=self.size,padding_reused_rows=extra,
	steps=self.steps_per_epoch,
	dataloader_repeats_per_optimizer_step=self.dataloader_repeats))+'\n')
	# TRL 0.21 consumes one 512-row dataloader batch for generation and
	# then 511 more dataloader batches while training the 512 buffered
	# microbatches. Those subsequent inputs are ignored. Repeat each
	# complete generation group so the next optimizer step advances to
	# the next group, rather than skipping 511 groups or ending an epoch.
	for start in range(0,len(order),self.prompts_per_step):
	group=order[start:start+self.prompts_per_step]
	for _ in range(self.dataloader_repeats):
	for index in group:
	for _ in range(self.generations):yield index