wehub-resource-sync
533 lines
25 KiB
Python
533 lines
25 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import inspect
|
|
import os
|
|
import random
|
|
import torch
|
|
import torch.nn as nn
|
|
import trl
|
|
from accelerate.utils import gather_object, is_peft_model
|
|
from collections import defaultdict, deque
|
|
from contextlib import contextmanager, nullcontext
|
|
from copy import deepcopy
|
|
from packaging import version
|
|
from transformers import PreTrainedModel, Trainer
|
|
from trl import SFTTrainer as HFSFTTrainer
|
|
from trl.trainer.utils import RepeatSampler
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
from swift.rl_core.data import GKDBatch, GKDSample
|
|
from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_output, build_teacher_requests, encode_gkd_samples,
|
|
fetch_teacher_parsed_by_routing)
|
|
from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss
|
|
from swift.template import TemplateInputs
|
|
from swift.trainers import SwiftMixin, disable_gradient_checkpointing
|
|
from swift.utils import (JsonlWriter, get_logger, is_swanlab_available, is_wandb_available, remove_response,
|
|
swanlab_get_run, to_device)
|
|
from .rollout_mixin import DataType, RolloutTrainerMixin
|
|
from .utils import get_gather_if_zero3_context, identity_data_collator, profiling_decorator
|
|
|
|
try:
|
|
from liger_kernel.chunked_loss import LigerFusedLinearJSDLoss
|
|
_liger_kernel_available = True
|
|
except ImportError:
|
|
_liger_kernel_available = False
|
|
|
|
if version.parse(trl.__version__) >= version.parse('0.26.0'):
|
|
from trl.experimental.gkd import GKDTrainer as HFGKDTrainer
|
|
else:
|
|
from trl import GKDTrainer as HFGKDTrainer
|
|
|
|
del HFGKDTrainer.__init__
|
|
del HFSFTTrainer.__init__
|
|
|
|
logger = get_logger()
|
|
if is_wandb_available():
|
|
import wandb
|
|
if is_swanlab_available():
|
|
import swanlab
|
|
|
|
|
|
class GKDTrainer(RolloutTrainerMixin, SwiftMixin, HFGKDTrainer):
|
|
|
|
sample_cls = GKDSample
|
|
|
|
def __init__(self, model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, *_args, **kwargs):
|
|
self.vllm_client = kwargs.pop('vllm_client', None)
|
|
self.gkd_logits_topk = kwargs.pop('gkd_logits_topk', None)
|
|
self._pop_teacher_kwargs(kwargs) # teacher kwargs (shared with GRPO via RolloutTrainerMixin)
|
|
super().__init__(model, None, *_args, **kwargs)
|
|
args = kwargs['args']
|
|
self.lmbda = args.lmbda
|
|
self.temperature = args.temperature
|
|
self._metrics = {'train': defaultdict(list), 'eval': defaultdict(list)}
|
|
self._total_train_tokens = 0
|
|
|
|
# Initialize logging components
|
|
self._prepare_logging()
|
|
|
|
# Initialize liger loss if enabled
|
|
self._prepare_liger_loss()
|
|
|
|
self._setup_teacher()
|
|
|
|
# Initialize rollout infrastructure for vLLM support
|
|
self.prepare_rollout()
|
|
|
|
# Initialize resample data iterator for truncation_strategy 'raise'('delete')
|
|
if self.template.truncation_strategy == 'raise':
|
|
self._prepare_resample_data_iterator()
|
|
|
|
self._step = 0
|
|
self._buffered_inputs = None
|
|
|
|
def _get_data_collator(self, args, template):
|
|
return identity_data_collator
|
|
|
|
def _get_train_sampler(self, train_dataset=None):
|
|
return RepeatSampler(
|
|
data_source=train_dataset or self.train_dataset,
|
|
mini_repeat_count=1,
|
|
batch_size=self.args.generation_batch_size,
|
|
repeat_count=self.args.steps_per_generation,
|
|
shuffle=getattr(self, 'shuffle_dataset', True),
|
|
seed=self.args.seed,
|
|
)
|
|
|
|
def get_train_dataloader(self):
|
|
return self._get_dataloader(
|
|
dataset=self.train_dataset,
|
|
description='Training',
|
|
batch_size=self._train_batch_size * self.args.steps_per_generation,
|
|
sampler_fn=self._get_train_sampler,
|
|
is_training=True,
|
|
)
|
|
|
|
def _compute_jsd_loss(self, student_logits, teacher_output: TeacherOutput, labels):
|
|
"""Compute JSD loss. teacher_output.labels is always set (equals student labels when non-OPSD)."""
|
|
shifted_labels = torch.roll(labels, shifts=-1, dims=1)
|
|
teacher_output.labels = torch.roll(teacher_output.labels, shifts=-1, dims=1)
|
|
if self.gkd_logits_topk is not None:
|
|
teacher_output = teacher_output.to_topk(self.gkd_logits_topk)
|
|
total, num_valid = gkd_loss(student_logits, teacher_output, shifted_labels, self.beta, self.temperature)
|
|
if num_valid == 0:
|
|
return total * 0
|
|
return total / num_valid
|
|
|
|
@profiling_decorator
|
|
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
|
|
model_inputs = inputs['model_inputs']
|
|
gkd_batch: GKDBatch = inputs['gkd_batch']
|
|
data_source = gkd_batch.data_source
|
|
teacher_model_inputs = inputs['teacher_model_inputs']
|
|
|
|
# logits_to_keep trims each forward to its completion region. student and teacher
|
|
# are masked independently in the loss via their own labels (extract_active), so the
|
|
# two sides can be trimmed independently even when OPSD makes their lengths differ.
|
|
use_logits_to_keep = self.get_use_logits_to_keep(self.template.sequence_parallel_size == 1)
|
|
if use_logits_to_keep and not self.use_liger_gkd_loss:
|
|
self.prepare_logits_to_keep(model_inputs)
|
|
# Trim the teacher only when it runs a forward. Teacher API does not (top-k is
|
|
# pre-assembled on the full sequence). For a local teacher forward, non-OPSD reuses
|
|
# the student encoding (separate dict) and OPSD has its own encoding — both are trimmed
|
|
# here so the teacher forward also benefits, while extract_active aligns each side via
|
|
# its own labels.
|
|
if not self.use_teacher_api:
|
|
self.prepare_logits_to_keep(teacher_model_inputs)
|
|
|
|
# Teacher labels come from the teacher-side encoding (== student encoding when non-OPSD).
|
|
# Must be read AFTER prepare_logits_to_keep which may truncate labels.
|
|
teacher_labels = teacher_model_inputs['labels']
|
|
|
|
# model_inputs is clean from template.encode; exclude 'labels' for model forward.
|
|
forward_inputs = {k: v for k, v in model_inputs.items() if k != 'labels'}
|
|
|
|
teacher_fwd_inputs = {k: v for k, v in teacher_model_inputs.items() if k != 'labels'}
|
|
|
|
if self.use_liger_gkd_loss:
|
|
# Liger fused JSD loss for memory efficiency
|
|
# Get base models (exclude lm_head to save memory)
|
|
unwrapped_student = self.accelerator.unwrap_model(model)
|
|
if is_peft_model(unwrapped_student):
|
|
unwrapped_student = unwrapped_student.base_model.model
|
|
base_student = getattr(unwrapped_student, getattr(unwrapped_student, 'base_model_prefix', 'model'),
|
|
unwrapped_student)
|
|
|
|
unwrapped_teacher = self.accelerator.unwrap_model(self.teacher_model)
|
|
base_teacher = getattr(unwrapped_teacher, getattr(unwrapped_teacher, 'base_model_prefix', 'model'),
|
|
unwrapped_teacher)
|
|
|
|
# Forward through base models
|
|
student_outputs = base_student(**forward_inputs, use_cache=False)
|
|
|
|
load_context = self.load_teacher_model_context() if self.args.offload_teacher_model else nullcontext()
|
|
with load_context:
|
|
with torch.no_grad(), disable_gradient_checkpointing(self.teacher_model,
|
|
self.args.gradient_checkpointing_kwargs):
|
|
teacher_outputs = base_teacher(**forward_inputs, use_cache=False)
|
|
|
|
# Get hidden states (shifted)
|
|
student_hidden = student_outputs.last_hidden_state[:, :-1]
|
|
teacher_hidden = teacher_outputs.last_hidden_state[:, :-1]
|
|
|
|
# Release full outputs to free memory
|
|
del student_outputs, teacher_outputs
|
|
|
|
# Prepare labels (shifted)
|
|
labels_mask = model_inputs['labels'] != -100
|
|
masked_input_ids = torch.where(labels_mask, model_inputs['input_ids'],
|
|
torch.full_like(model_inputs['input_ids'], -100))
|
|
true_labels = masked_input_ids[:, 1:].contiguous()
|
|
|
|
# Release intermediate tensors
|
|
del labels_mask, masked_input_ids
|
|
|
|
# Get output heads
|
|
student_head = unwrapped_student.get_output_embeddings()
|
|
teacher_head = unwrapped_teacher.get_output_embeddings()
|
|
|
|
# Prepare context managers for gathering parameters in zero3
|
|
teacher_context = get_gather_if_zero3_context(self, is_zero3=self.is_teacher_ds3)(teacher_head.weight)
|
|
student_context = get_gather_if_zero3_context(self)(student_head.weight)
|
|
|
|
with teacher_context, student_context:
|
|
# Compute liger fused JSD loss
|
|
loss = self.liger_jsd_loss(
|
|
student_input=student_hidden,
|
|
student_weight=student_head.weight,
|
|
teacher_input=teacher_hidden,
|
|
teacher_weight=teacher_head.weight,
|
|
true_labels=true_labels,
|
|
student_bias=getattr(student_head, 'bias', None),
|
|
teacher_bias=getattr(teacher_head, 'bias', None),
|
|
)
|
|
# Release hidden states after loss computation
|
|
del student_hidden, teacher_hidden, true_labels
|
|
outputs_student = None
|
|
# Non-liger path: student forward + teacher output construction + JSD loss
|
|
else:
|
|
if self.args.sft_alpha > 0:
|
|
forward_inputs['labels'] = model_inputs['labels']
|
|
outputs_student = model(**forward_inputs)
|
|
|
|
# Construct teacher_out — the only part that differs across modes
|
|
if self.use_teacher_api:
|
|
teacher_topk_logprobs = gkd_batch.teacher_topk_logprobs
|
|
teacher_topk_indices = gkd_batch.teacher_topk_indices
|
|
assert teacher_topk_logprobs is not None and teacher_topk_indices is not None, (
|
|
'teacher_topk_logprobs/teacher_topk_indices missing on gkd_batch; '
|
|
'ensure _prepare_inputs (or prediction_step) populated them.')
|
|
teacher_out = TeacherOutput(
|
|
topk_logprobs=teacher_topk_logprobs,
|
|
topk_indices=teacher_topk_indices,
|
|
labels=teacher_labels,
|
|
)
|
|
else:
|
|
t_fwd = teacher_fwd_inputs
|
|
if self._is_self_distillation:
|
|
adapter_ctx = (
|
|
self.accelerator.unwrap_model(model).disable_adapter()
|
|
if self._teacher_use_disable_adapter else nullcontext())
|
|
with torch.no_grad(), adapter_ctx, \
|
|
disable_gradient_checkpointing(model, self.args.gradient_checkpointing_kwargs):
|
|
outputs_teacher = model(**t_fwd)
|
|
else:
|
|
assert self.teacher_model is not None
|
|
load_context = (
|
|
self.load_teacher_model_context() if self.args.offload_teacher_model else nullcontext())
|
|
with torch.no_grad(), load_context, disable_gradient_checkpointing(
|
|
self.teacher_model, self.args.gradient_checkpointing_kwargs):
|
|
outputs_teacher = self.teacher_model(**t_fwd)
|
|
teacher_out = TeacherOutput(full_logits=outputs_teacher.logits, labels=teacher_labels)
|
|
|
|
loss = self._compute_jsd_loss(outputs_student.logits, teacher_out, model_inputs['labels'])
|
|
|
|
if self.args.sft_alpha > 0 and data_source != DataSource.STUDENT:
|
|
loss = loss + self.args.sft_alpha * outputs_student.loss
|
|
|
|
# Return loss
|
|
if return_outputs:
|
|
return (loss, outputs_student)
|
|
else:
|
|
return loss
|
|
|
|
def _log_student_completions(self, generated_samples) -> None:
|
|
"""Log student completions from generated GKDSample list."""
|
|
if not self.log_completions:
|
|
return
|
|
messages = [s.messages[:-1] for s in generated_samples]
|
|
completions = [deepcopy(s.messages[-1]['content']) for s in generated_samples]
|
|
valid_messages = gather_object(messages)
|
|
valid_completions = gather_object(completions)
|
|
self._logs['prompt'].extend(self._apply_chat_template_to_messages_list(valid_messages))
|
|
self._logs['completion'].extend(valid_completions)
|
|
|
|
def _encode_samples(self, samples: List[GKDSample]):
|
|
"""Encode student + teacher (OPSD-aware). Mirrors Megatron GKD _encode_samples.
|
|
|
|
Args:
|
|
samples: Samples to encode (may be generated samples for STUDENT mode).
|
|
|
|
Returns:
|
|
(encoded_inputs, teacher_encoded_or_None) — ``teacher_encoded`` is non-None
|
|
only for OPSD; non-OPSD returns None (teacher reuses the student encoding).
|
|
"""
|
|
template = self.template
|
|
|
|
with self._template_context(template):
|
|
student_encoded_list, teacher_encoded_list, has_opsd = encode_gkd_samples(samples, template)
|
|
encoded_inputs = to_device(template.data_collator(student_encoded_list), self.model.device)
|
|
if has_opsd:
|
|
with self._template_context(template):
|
|
teacher_encoded = to_device(template.data_collator(teacher_encoded_list), self.model.device)
|
|
else:
|
|
teacher_encoded = None
|
|
|
|
return encoded_inputs, teacher_encoded
|
|
|
|
def _build_teacher_requests(self, samples: List[GKDSample]):
|
|
"""Build teacher API requests from samples (mirrors Megatron GKD)."""
|
|
return build_teacher_requests(samples, self.template)
|
|
|
|
def _rollout_samples(self, inputs: DataType) -> List[GKDSample]:
|
|
"""Pick the data source, convert rows to samples, and generate for the STUDENT branch."""
|
|
if self._get_random_num() <= self.lmbda:
|
|
self._data_source = DataSource.STUDENT
|
|
else:
|
|
self._data_source = DataSource.DATASET
|
|
|
|
if self.template.truncation_strategy == 'raise':
|
|
inputs = self.resample_encode_failed_inputs(inputs, strip_response=self._data_source != DataSource.DATASET)
|
|
samples = self.to_samples(inputs)
|
|
|
|
if self._data_source == DataSource.STUDENT:
|
|
# Generation at full-batch level (uses _generate_completions from mixin)
|
|
samples = self._generate_completions(samples)
|
|
self._log_student_completions(samples)
|
|
return samples
|
|
|
|
def _score_completions(self, samples: List[GKDSample]) -> List[GKDSample]:
|
|
"""GKD has no reward scoring; teacher signals are filled in ``_postprocess_batch``."""
|
|
return samples
|
|
|
|
@profiling_decorator
|
|
def _prepare_batch_inputs(self, samples: List[GKDSample]) -> List[Dict[str, Any]]:
|
|
"""Encode + collate samples into per-micro-batch ``{model_inputs, gkd_batch}`` dicts.
|
|
|
|
Mirrors GRPO ``_prepare_batch_inputs``: takes all samples and splits internally.
|
|
"""
|
|
sample_chunks = self.split_by_mini_batches(samples)
|
|
ga_batch_encoded_inputs: List[Dict[str, Any]] = []
|
|
for chunk_samples in sample_chunks:
|
|
encoded_inputs, teacher_encoded = self._encode_samples(chunk_samples)
|
|
# teacher_encoded is non-None only for OPSD. Non-OPSD reuses the student encoding
|
|
# for the teacher, but as a separate dict so each side can run prepare_logits_to_keep
|
|
# independently without aliasing.
|
|
teacher_model_inputs = teacher_encoded or dict(encoded_inputs)
|
|
result: Dict[str, Any] = {
|
|
'model_inputs': encoded_inputs,
|
|
'gkd_batch': GKDBatch(data_source=self._data_source),
|
|
'teacher_model_inputs': teacher_model_inputs,
|
|
}
|
|
if self.use_teacher_api:
|
|
result['_teacher_requests'] = self._build_teacher_requests(chunk_samples)
|
|
result['_chunk_samples'] = chunk_samples
|
|
ga_batch_encoded_inputs.append(result)
|
|
return ga_batch_encoded_inputs
|
|
|
|
def _postprocess_batch(self, samples: List[GKDSample], batch_encoded_inputs: List[Dict[str, Any]]) -> None:
|
|
"""Fetch teacher logprobs via API and fill them into each ``gkd_batch``."""
|
|
if self.use_teacher_api:
|
|
self._fetch_and_assemble_teacher_logprobs(batch_encoded_inputs)
|
|
|
|
def _log_rollout(self, samples: List[GKDSample]) -> None:
|
|
"""Student completions are logged in ``_rollout_samples``; nothing extra here."""
|
|
|
|
@profiling_decorator
|
|
def _prepare_inputs(self, inputs: DataType) -> Dict[str, torch.Tensor]:
|
|
mode = 'train' if self.model.training else 'eval'
|
|
steps_per_generation = self.args.steps_per_generation
|
|
if mode == 'train':
|
|
if self._step % steps_per_generation == 0 or self._buffered_inputs is None:
|
|
self._buffered_inputs = self._generate_and_score_completions(inputs)
|
|
inputs = self._buffered_inputs[self._step % steps_per_generation]
|
|
self._step += 1
|
|
else:
|
|
inputs = self._generate_and_score_completions(inputs)[0]
|
|
return inputs
|
|
|
|
def _fetch_and_assemble_teacher_logprobs(self, chunks):
|
|
"""Fetch teacher logprobs via API using sample-based _teacher_requests.
|
|
|
|
Each sample routes to exactly one teacher by tag (single teacher = all samples). Routing
|
|
runs once over the flattened chunks, so a single teacher is one DP gather → infer → slice.
|
|
The parsed results are sliced back per chunk and fed to ``assemble_teacher_output``.
|
|
"""
|
|
flat_samples, flat_requests, chunk_sizes = [], [], []
|
|
for c in chunks:
|
|
reqs = c.pop('_teacher_requests', [])
|
|
flat_samples.extend(c.pop('_chunk_samples', []))
|
|
flat_requests.extend(reqs)
|
|
chunk_sizes.append(len(reqs))
|
|
|
|
parsed = fetch_teacher_parsed_by_routing(
|
|
flat_samples,
|
|
flat_requests,
|
|
self.teacher_configs,
|
|
self.teacher_clients,
|
|
gather_fn=self._gather_teacher_requests,
|
|
infer_fn=lambda handle, client: self._infer_teacher_requests(
|
|
handle, topk=self.gkd_logits_topk, teacher_client=client),
|
|
scatter_fn=self._scatter_teacher_parsed,
|
|
is_main_process=self.accelerator.is_main_process,
|
|
tag_key=self.args.teacher_tag_key)
|
|
|
|
per_chunk_parsed, offset = [], 0
|
|
for cs in chunk_sizes:
|
|
per_chunk_parsed.append(parsed[offset:offset + cs])
|
|
offset += cs
|
|
|
|
for c, chunk_parsed in zip(chunks, per_chunk_parsed):
|
|
gkd_batch: GKDBatch = c['gkd_batch']
|
|
target = c['teacher_model_inputs']
|
|
teacher_out = assemble_teacher_output(
|
|
chunk_parsed,
|
|
teacher_model_inputs=target,
|
|
topk=self.gkd_logits_topk,
|
|
template_padding_free=self.template.padding_free,
|
|
device=target['input_ids'].device,
|
|
)
|
|
gkd_batch.teacher_topk_logprobs = teacher_out.topk_logprobs
|
|
gkd_batch.teacher_topk_indices = teacher_out.topk_indices
|
|
|
|
@profiling_decorator
|
|
def training_step(self,
|
|
model: nn.Module,
|
|
inputs: DataType,
|
|
num_items_in_batch: Optional[int] = None) -> torch.Tensor:
|
|
return HFSFTTrainer.training_step(self, model, inputs, num_items_in_batch)
|
|
|
|
def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
|
|
inputs = self._prepare_inputs(inputs)
|
|
with torch.no_grad():
|
|
with self.compute_loss_context_manager():
|
|
loss = self.compute_loss(model, inputs)
|
|
loss = loss.mean().detach()
|
|
return loss, None, None
|
|
|
|
@contextmanager
|
|
def offload_context(self):
|
|
"""Offload student model and optimizer to CPU during vLLM on-policy generation."""
|
|
if self.args.offload_model:
|
|
self.offload_model(self.accelerator.unwrap_model(self.model))
|
|
if getattr(self, 'optimizer', None) and self.args.offload_optimizer:
|
|
self.offload_optimizer()
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
if self.args.offload_model:
|
|
self.load_model(self.accelerator.unwrap_model(self.model))
|
|
if getattr(self, 'optimizer', None) and self.args.offload_optimizer:
|
|
self.load_optimizer()
|
|
|
|
def _get_random_num(self) -> float:
|
|
"""
|
|
Generate a deterministic random number.
|
|
|
|
Uses an isolated Random instance to avoid interfering with the global
|
|
random state, ensuring thread-safety and consistent behavior across processes.
|
|
|
|
Returns:
|
|
float: A random number in the range [0.0, 1.0).
|
|
"""
|
|
seed = int(getattr(self.args, 'seed', 0))
|
|
seed += int(self.state.global_step)
|
|
rng = random.Random(seed)
|
|
return rng.random()
|
|
|
|
def _prepare_liger_loss(self):
|
|
"""Initialize liger loss if enabled."""
|
|
args = self.args
|
|
self.use_liger_gkd_loss = False
|
|
if getattr(args, 'use_liger_kernel', False):
|
|
if not _liger_kernel_available:
|
|
raise ImportError(
|
|
'Liger kernel is not installed. Please install liger-kernel by running: pip install liger-kernel')
|
|
assert self.args.sft_alpha == 0, 'SFT loss is not supported with liger loss'
|
|
assert self.gkd_logits_topk is None, 'Top-k mode is not supported with liger loss'
|
|
self.liger_jsd_loss = LigerFusedLinearJSDLoss(
|
|
beta=self.beta,
|
|
ignore_index=-100,
|
|
temperature=self.temperature,
|
|
compiled=False,
|
|
)
|
|
self.use_liger_gkd_loss = True
|
|
|
|
def _prepare_logging(self):
|
|
"""Initialize logging components for on-policy rollout tracking."""
|
|
args = self.args
|
|
self.log_completions = args.log_completions
|
|
self.wandb_log_unique_prompts = getattr(args, 'wandb_log_unique_prompts', False)
|
|
self.jsonl_writer = JsonlWriter(os.path.join(self.args.output_dir, 'completions.jsonl'))
|
|
|
|
# Initialize logs deque for storing rollout data (aligned with GRPO)
|
|
self._logs = {
|
|
'prompt': deque(),
|
|
'completion': deque(),
|
|
}
|
|
|
|
def _apply_chat_template_to_messages_list(self, messages_list: DataType):
|
|
"""Convert messages list to prompt text list using template (aligned with GRPO)."""
|
|
prompts_text = []
|
|
for messages in messages_list:
|
|
remove_response(messages)
|
|
template_inputs = TemplateInputs.from_dict({'messages': messages})
|
|
res = self.template.encode(template_inputs)
|
|
prompts_text.append(self.template.safe_decode(res['input_ids']))
|
|
return prompts_text
|
|
|
|
def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None:
|
|
"""Override log method to include completion table logging (aligned with GRPO)."""
|
|
# Call parent log method
|
|
import transformers
|
|
from packaging import version
|
|
if version.parse(transformers.__version__) >= version.parse('4.47.0.dev0'):
|
|
super().log(logs, start_time)
|
|
else:
|
|
super().log(logs)
|
|
|
|
# Log completions table if we have data (only for on-policy generations)
|
|
if self.accelerator.is_main_process and self.log_completions and len(self._logs['prompt']) > 0:
|
|
seen_nums = len(self._logs['prompt'])
|
|
table = {
|
|
'step': [str(self.state.global_step)] * seen_nums,
|
|
'prompt': list(self._logs['prompt'])[:seen_nums],
|
|
'completion': list(self._logs['completion'])[:seen_nums],
|
|
}
|
|
|
|
# Write to jsonl
|
|
self.jsonl_writer.append(table)
|
|
|
|
self._logs['prompt'].clear()
|
|
self._logs['completion'].clear()
|
|
# Log to wandb if enabled
|
|
report_to_wandb = self.args.report_to and 'wandb' in self.args.report_to and wandb.run is not None
|
|
if report_to_wandb:
|
|
wandb_table = table.copy()
|
|
import pandas as pd
|
|
df = pd.DataFrame(wandb_table)
|
|
if self.wandb_log_unique_prompts:
|
|
df = df.drop_duplicates(subset=['prompt'])
|
|
wandb.log({'completions': wandb.Table(dataframe=df)})
|
|
|
|
# Log to swanlab if enabled
|
|
report_to_swanlab = self.args.report_to and 'swanlab' in self.args.report_to and swanlab_get_run(
|
|
) is not None
|
|
if report_to_swanlab:
|
|
headers = list(table.keys())
|
|
rows = []
|
|
for i in range(len(table['step'])):
|
|
row = [table[header][i] for header in headers]
|
|
rows.append(row)
|
|
swanlab.log({'completions': swanlab.echarts.Table().add(headers, rows)})
|