project-llm-trainer 0.3__py3-none-any.whl

llm_trainer/utils.py

@@ -0,0 +1,262 @@
+import random
+from typing import Tuple, Optional
+import torch
+from torch.nn.utils.rnn import pad_sequence
+import torch.nn.functional as F
+from .tools import TrainerTools
+import numpy as np
+
+
+def set_seed(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def extra_image_tag_and_repeat_image_tok(
+        inputs: list[int],
+        tokens_per_image: int
+) -> Tuple[list[int], Optional[int]]:
+    # tokens_per_image=3 -> <image>{image_tag}...xxxx -> <image><image><image>...xxx
+    image_tok = TrainerTools().tokenizer.image
+    if image_tok not in inputs:
+        return inputs, None
+
+    image_tok_idx = inputs.index(image_tok)
+    image_tag_idx = image_tok_idx + 1
+
+    if image_tag_idx < len(inputs):
+        # remove it
+        image_tag = inputs.pop(image_tag_idx)
+    else:
+        image_tag = None
+
+    # repeat image_tok
+    new_inputs = inputs[:image_tok_idx] + [image_tok] * tokens_per_image + inputs[image_tok_idx + 1:]
+    return new_inputs, image_tag
+
+
+def batch_extra_image_tag_and_repeat_image_tok(
+        tokens: torch.Tensor,
+        tokens_per_image: int
+) -> Tuple[torch.Tensor, list[int]]:
+    new_tokens = []
+    image_tags = []
+
+    tokens_list = tokens.cpu().detach().tolist()
+    for token in tokens_list:
+        new_token, image_tag = extra_image_tag_and_repeat_image_tok(token, tokens_per_image)
+        new_tokens.append(new_token)
+        image_tags.append(image_tag)
+
+    return torch.tensor(new_tokens, dtype=tokens.dtype, device=tokens.device), image_tags
+
+
+def repeat_image_tok(
+        tokens: torch.Tensor,
+        tokens_per_image: int
+) -> torch.Tensor:
+    # tokens_per_image=3 -> <image>...xxxx -> <image><image><image>...xxx
+    image_tok = TrainerTools().tokenizer.image
+    if image_tok not in tokens:
+        return tokens
+
+    image_tok_idx = torch.where(tokens == image_tok)[0].item()
+    repeat_image_toks = torch.tensor([image_tok] * tokens_per_image, dtype=tokens.dtype, device=tokens.device)
+
+    # repeat image_tok
+    new_tokens = torch.concat([tokens[:image_tok_idx], repeat_image_toks, tokens[image_tok_idx + 1:]], dim=-1)
+    return new_tokens
+
+
+def batch_repeat_image_tok(
+        tokens: torch.Tensor,
+        tokens_per_image: int
+) -> torch.Tensor:
+    new_tokens = []
+
+    for token in tokens:
+        new_tokens.append(repeat_image_tok(token, tokens_per_image))
+
+    return torch.stack(new_tokens, dim=0)
+
+
+def _pad_sequence(batch_data):
+    # [[x,x,x], [y,y,y]]
+    inputs = pad_sequence(batch_data, batch_first=True, padding_value=TrainerTools().tokenizer.pad)
+    # crossEntropy默认的ignore_index是-100
+    labels = pad_sequence(batch_data, batch_first=True, padding_value=-100)
+
+    return inputs, labels
+
+
+def _mask_prompt(labels):
+    tokenizer = TrainerTools().tokenizer
+    # 支持多轮会话的mask
+    for batch, label in enumerate(labels):
+        start_index = -1
+        for index, token in enumerate(label):
+            if token == tokenizer.system or token == tokenizer.user:
+                start_index = index
+            elif token == tokenizer.end and start_index != -1:
+                labels[batch, start_index:index + 1] = -100
+                start_index = -1
+
+    return labels
+
+
+def _zero_pad_sequences(
+    sequences: list[torch.Tensor], side: str = "left"
+) -> torch.Tensor:
+    assert side in ("left", "right")
+    max_len = max(seq.size(0) for seq in sequences)
+    padded_sequences = []
+    for seq in sequences:
+        pad_len = max_len - seq.size(0)
+        padding = (pad_len, 0) if side == "left" else (0, pad_len)
+        padded_sequences.append(F.pad(seq, padding))
+    return torch.stack(padded_sequences, dim=0)
+
+
+def pretrain_collate_fn(batch_data):
+    inputs, labels = _pad_sequence(batch_data)
+
+    # inputs, labels
+    return {'inputs': inputs, 'labels': labels}
+
+
+def get_sft_collate_fn(mask_prompt: bool):
+    def sft_collate_fn(batch_data):
+        """
+         如果是sft，则不计算prompt部分的loss, 例如：
+        logits: [USER]你好[BOT]我好[SEP]
+        labels: [USER]你好[BOT]我好[SEP]
+
+        shift_logits: [USER]你好[BOT]我好
+        shift_labels: 你好[BOT]我好[SEP]
+
+        mask_labels: mask mask mask mask 我好[SEP]
+            * mask=-100和pad一样
+
+
+        多轮对话场景
+        [USER]你好[BOT]我好[SEP][USER]很好[BOT]不好[SEP]
+        mask: mask mask mask mask 我好[SEP] mask mask mask mask 不好[SEP]
+        """
+        batch_train_data = []
+        image_tags = []
+        for item in batch_data:
+            batch_train_data.append(item['inputs'])
+            image_tags.append(item['image_tag'])
+
+        inputs, labels = _pad_sequence(batch_train_data)
+        if mask_prompt:
+            labels = _mask_prompt(labels)
+
+        return {'inputs': inputs, 'labels': labels, 'image_tags': image_tags}
+
+    return sft_collate_fn
+
+
+def get_dpo_collate_fn(mask_prompt: bool):
+    def dpo_collate_fn(batch_data):
+        # batch_data: [{'chosen': chosen, 'rejected': rejected}, {'chosen': chosen, 'rejected': rejected}]
+        chosen_inputs = []
+        chosen_labels = []
+        rejected_inputs = []
+        rejected_labels = []
+
+        max_len = 0
+        for key in ['chosen', 'rejected']:
+            max_len = max(max(len(item[key]) for item in batch_data), max_len)
+
+        for item in batch_data:
+            chosen_sequence = item['chosen']
+            chosen_inputs.append(chosen_sequence + [TrainerTools().tokenizer.pad] * (max_len - len(chosen_sequence)))
+            chosen_labels.append(chosen_sequence + [-100] * (max_len - len(chosen_sequence)))
+
+            rejected_sequence = item['rejected']
+            rejected_inputs.append(rejected_sequence + [TrainerTools().tokenizer.pad] * (max_len - len(rejected_sequence)))
+            rejected_labels.append(rejected_sequence + [-100] * (max_len - len(rejected_sequence)))
+
+        chosen_inputs = torch.tensor(chosen_inputs).long()
+        chosen_labels = torch.tensor(chosen_labels).long()
+        if mask_prompt:
+            chosen_labels = _mask_prompt(chosen_labels)
+
+        rejected_inputs = torch.tensor(rejected_inputs).long()
+        rejected_labels = torch.tensor(rejected_labels).long()
+        if mask_prompt:
+            rejected_labels = _mask_prompt(rejected_labels)
+
+        return {
+            'chosen_inputs': chosen_inputs,
+            'chosen_labels': chosen_labels,
+            'rejected_inputs': rejected_inputs,
+            'rejected_labels': rejected_labels
+        }
+
+    return dpo_collate_fn
+
+
+def split_batch(data_per_batch: dict) -> list[dict]:
+    """
+    from: data_per_batch("sequences": [group_size, max_generate_len] ...)
+    to:   [dict("sequences": [max_generate_len] ...) ... group_size]
+    """
+
+    group_size = data_per_batch['sequence_ids'].size(0)
+    # [{"sequence_ids": xxx, "old_log_probs": xxx...}, ...]
+    group_data = [{} for _ in range(group_size)]
+
+    keys = (
+        'sequence_ids',
+        'old_log_probs',
+        'ref_log_probs',
+        'advantages',
+        'attention_mask',
+        'mask',
+    )
+
+    for key in keys:
+        value = data_per_batch[key]
+        if value is None:
+            vals = [None] * group_size
+        else:
+            vals = torch.unbind(value)
+
+        for i, v in enumerate(vals):
+            group_data[i][key] = v
+
+    return group_data
+
+
+def join_batch(batch_data: list[dict]) -> dict:
+    """
+    from: [dict("sequences": [max_generate_len] ...), ...]
+    to:   dict("sequences": max_generate_len, ...)
+    """
+
+    result = {}
+    keys = (
+        'sequence_ids',
+        'old_log_probs',
+        'ref_log_probs',
+        'advantages',
+        'attention_mask',
+        'mask',
+    )
+
+    for key in keys:
+        # [sequence_ids, sequence_ids ...]
+        # shape [batch_size, seq_len]
+        vals = [item[key] for item in batch_data]
+        if all(v is not None for v in vals):
+            data = _zero_pad_sequences(vals, "left")
+        else:
+            data = None
+        result[key] = data
+
+    return result

project_llm_trainer-0.3.data/scripts/calc_intermediate_size

@@ -0,0 +1,15 @@
+#!python
+
+if __name__ == '__main__':
+    import sys
+    arguments = sys.argv[1:]
+    hidden_size = int(arguments[0])
+    if len(arguments) > 1:
+        multiple_of = int(arguments[1])
+    else:
+        multiple_of = 64
+
+    intermediate_size = 4 * hidden_size
+    intermediate_size = int(2 * intermediate_size / 3)
+    intermediate_size = multiple_of * ((intermediate_size + multiple_of - 1) // multiple_of)
+    print(f'intermediate_size={intermediate_size}')

project_llm_trainer-0.3.data/scripts/ddp_train

@@ -0,0 +1,12 @@
+#!python
+
+if __name__ == '__main__':
+    import os, sys
+    arguments = sys.argv[1:]
+    run_file_name = arguments[0]
+
+    os.environ['PARALLEL_TYPE'] = 'ddp'
+    command = f'torchrun --standalone --nproc_per_node=gpu {run_file_name}'
+
+    print(f'real command is {command}')
+    os.system(command)

project_llm_trainer-0.3.data/scripts/ds_train

@@ -0,0 +1,12 @@
+#!python
+
+if __name__ == '__main__':
+    import os, sys
+    arguments = sys.argv[1:]
+    run_file_name = arguments[0]
+
+    os.environ['PARALLEL_TYPE'] = 'ds'
+    command = f'deepspeed {run_file_name}'
+
+    print(f'real command is {command}')
+    os.system(command)

project_llm_trainer-0.3.data/scripts/plot_loss

@@ -0,0 +1,39 @@
+#!python
+
+if __name__ == '__main__':
+    import os, sys
+    import matplotlib.pyplot as plt
+
+    arguments = sys.argv[1:]
+    loss_file = arguments[0]
+
+    if not os.path.exists(loss_file):
+        print(f'{loss_file} not found')
+        exit(0)
+
+    steps = []
+    losses = []
+    with open(loss_file, 'r') as f:
+        step = 0
+        for line in f:
+            if not line or 'loss:' not in line:
+                if 'start train' not in line:
+                    steps.clear()
+                    losses.clear()
+                    step = 0
+                continue
+
+            # (2025-03-19 20:13:44) epoch: 0, file: 1/1, batch: 623/1099, loss: 0.12186837196350098
+            loss = float(line.split('loss:')[-1].strip())
+
+            steps.append(step)
+            losses.append(loss)
+            step += 1
+
+    plt.xlabel('steps')
+    plt.ylabel('loss')
+
+    plt.plot(steps, losses)
+    plt.show()
+
+

project_llm_trainer-0.3.data/scripts/plot_lr

@@ -0,0 +1,41 @@
+#!python
+
+if __name__ == '__main__':
+    import os, sys
+    import matplotlib.pyplot as plt
+
+    arguments = sys.argv[1:]
+    lr_file = arguments[0]
+
+    if not os.path.exists(lr_file):
+        print(f'{lr_file} not found')
+        exit(0)
+
+    steps = []
+    lrs = []
+    with open(lr_file, 'r') as f:
+        for line in f:
+            if not line:
+                continue
+            # line: (2025-03-19 18:15:30) step=159,lr=2.159680442248444e-05
+            # data: 159,lr=2.159680442248444e-05
+            data = line.split('step=')[-1]
+            # [159, 2.159680442248444e-05]
+            data = data.split(',lr=')
+
+            step = int(data[0])
+            lr = float(data[1])
+
+            if step in steps:
+                continue
+
+            steps.append(step)
+            lrs.append(lr)
+
+    plt.xlabel('steps')
+    plt.ylabel('lr')
+
+    plt.plot(steps, lrs)
+    plt.show()
+
+

project_llm_trainer-0.3.data/scripts/py_train

@@ -0,0 +1,12 @@
+#!python
+
+if __name__ == '__main__':
+    import os, sys
+    arguments = sys.argv[1:]
+    run_file_name = arguments[0]
+
+    os.environ['PARALLEL_TYPE'] = 'none'
+    command = f'python3 {run_file_name}'
+
+    print(f'real command is {command}')
+    os.system(command)

project_llm_trainer-0.3.data/scripts/smart_train

@@ -0,0 +1,28 @@
+#!python
+
+if __name__ == '__main__':
+    import os, sys, torch
+    arguments = sys.argv[1:]
+    run_file_name = arguments[0]
+
+    try:
+        import deepspeed
+        parallel_type = 'ds'
+    except:
+        gpu_count = torch.cuda.device_count()
+        if gpu_count <= 1:
+            parallel_type = 'none'
+        else:
+            parallel_type = 'ddp'
+
+    os.environ['PARALLEL_TYPE'] = parallel_type
+
+    if parallel_type == 'ds':
+        command = f'deepspeed {run_file_name}'
+    elif parallel_type == 'ddp':
+        command = f'torchrun --standalone --nproc_per_node=gpu {run_file_name}'
+    else:
+        command = f'python3 {run_file_name}'
+
+    print(f'real command is {command}')
+    os.system(command)

project_llm_trainer-0.3.dist-info/METADATA

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: project_llm_trainer
+Version: 0.3
+Summary: LLM and VLM trainer
+Author: qibin
+Author-email: qibin0506@gmail.com
+Dynamic: author
+Dynamic: author-email
+Dynamic: summary

project_llm_trainer-0.3.dist-info/RECORD

@@ -0,0 +1,34 @@
+llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
+llm_trainer/checkpoint.py,sha256=iTbnmVrT0Ql4DpD178UI95zCmfBUdYtoJS5wIvf8_4k,6099
+llm_trainer/dataset.py,sha256=uz1TTd87ikf7CZPdGxmR95TSQTFWPPTilgWLBWO46_I,3916
+llm_trainer/dcp.py,sha256=PkD97DyrOtoTKn4FJsfL3VqAy4dxufgjdzJEz8-Cnoc,3635
+llm_trainer/dpo_trainer.py,sha256=6rm8Jq0rI0xazcl_bCOun8rnd34Tb_PKgezowhwoiCM,13150
+llm_trainer/ds_checkpoint.py,sha256=_svpzqRaa43--DKPputoXAelc6X9vPM0gNQu-hlh6NI,2153
+llm_trainer/eval.py,sha256=sCvdYnqWWf5_nuDQN5BHb_YivXLOQW-V0ET9mPu0tPU,2389
+llm_trainer/generate_utils.py,sha256=4iM0vyc_1C_iTL31GlS9PR4eZtYaELPRZ02KDSPZA9U,15158
+llm_trainer/grpo_trainer.py,sha256=gWDX8vRZ7hLKl_483X5ua92nst1m617BrqnzLhwr87g,16390
+llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
+llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
+llm_trainer/parallel.py,sha256=2VJtW3Gq2c1yS_LdcrNhk7B12prFwBmFnKhvV8FS2d8,4428
+llm_trainer/parallel_ddp.py,sha256=Gz-3LZ6LKmqlNwxrnGRC4uKoqoSxCvp9JHejIBSQp3c,1238
+llm_trainer/parallel_ds.py,sha256=W_PkczyAlgffCRcQadN-Pf7H7HM7TU26v5W63jKELFM,990
+llm_trainer/parallel_fsdp.py,sha256=u9XbbVTzcsMcaf-aQFrC_QwWsDRGoEpRmgvu1cKNtgk,3887
+llm_trainer/parallel_none.py,sha256=a6tt3aBmCq5rSP7n2I-sF-hsZ992BbLbpbxutDCFJfs,607
+llm_trainer/scheduler.py,sha256=Xz8HhwoRMjRe41sf_NHhpZfkTlEs0I2MYusvMY6hCVw,3531
+llm_trainer/sft_trainer.py,sha256=T9CujoEp8D5I65fLF2wgV6SPjzhGFbAI4We5NwL4O-M,1443
+llm_trainer/tokenizer.py,sha256=A7TYYUbtPf75kjCvWP7yBui4xZBObMk2aPem62YpwpY,6776
+llm_trainer/tools.py,sha256=AhfjN9oln5Pyif1SgCWwgQg-Q5acTCd9xpz4L26QUjA,3039
+llm_trainer/train_configs.py,sha256=FAlylSYVeh_oJGTy2fcMNUV8JLD6B70hMuk-iKx14iI,15748
+llm_trainer/trainer.py,sha256=mq51d-2ADUpcWCArszhYnOSTveatt3_x43hcC7IZgYk,24330
+llm_trainer/utils.py,sha256=04XiMENVotNgbNRBn9wadHu-cJHPxj0Xq-zzLJmNgZQ,8062
+project_llm_trainer-0.3.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.3.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.3.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.3.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.3.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.3.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.3.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.3.dist-info/METADATA,sha256=P64NiFbJzSd4QkFJ5udQ4qMyHUorPp3ex4F3eIdtVdU,193
+project_llm_trainer-0.3.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.3.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.3.dist-info/RECORD,,

project_llm_trainer-0.3.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.7.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

project_llm_trainer-0.3.dist-info/top_level.txt

@@ -0,0 +1 @@
+llm_trainer