Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+RM-EN-01-30-2026/data/rm_train.jsonl filter=lfs diff=lfs merge=lfs -text
+RM-EN-01-30-2026/model/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+SFT-EN-01-29-2026/data/train.jsonl filter=lfs diff=lfs merge=lfs -text
+SFT-EN-01-29-2026/model/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+sft_model_backup/tokenizer.json filter=lfs diff=lfs merge=lfs -text

RM-EN-01-30-2026/code/main.py

@@ -0,0 +1,394 @@
+#!/usr/bin/env python
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+#!/usr/bin/env python
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+#!/usr/bin/env python
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+import argparse
+import os
+import math
+import sys
+
+import torch
+from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
+from torch.utils.data.distributed import DistributedSampler
+
+from transformers import (
+    SchedulerType,
+    get_scheduler,
+)
+
+import deepspeed
+from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
+from deepspeed.accelerator import get_accelerator
+
+from dschat.utils.model.model_utils import create_critic_model
+from dschat.utils.data.data_utils import create_prompt_dataset, DataCollatorReward
+from dschat.utils.utils import print_rank_0, to_device, save_hf_format, save_hf_format_safetensors, set_random_seed, get_all_reduce_mean, get_optimizer_grouped_parameters, save_zero_three_model, load_hf_tokenizer
+from dschat.utils.ds_utils import get_train_ds_config
+from dschat.utils.module.lora import convert_linear_layer_to_lora, convert_lora_to_linear_layer, only_optimize_lora_parameters, make_model_gradient_checkpointing_compatible
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description=
+        "Finetune a transformers model on a causal language modeling task")
+    parser.add_argument('--data_path',
+                       nargs='*',
+                       default=['Dahoas/rm-static'],
+                       help='Path to the training dataset. Accepted format:'
+                       '1) a single data path, 2) multiple datasets in the'
+                       'form: dataset1-path dataset2-path ...')
+    parser.add_argument('--data_split',
+                       type=str,
+                       default='2,4,4',
+                       help='Comma-separated list of proportions for training'
+                       'phase 1, 2, and 3 data. For example the split `6,2,2`'
+                       'will use 60%% of data for phase 1, 20%% for phase 2'
+                       'and 20%% for phase 3.')
+    parser.add_argument(
+        '--data_output_path',
+        type=str,
+        default='/tmp/data_files/',
+        help=
+        'Where to store the data-related files such as shuffle index. This needs to be on a local storage of a node (not on a shared storage)'
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help=
+        "Path to pretrained model or model identifier from huggingface.co/models.",
+        required=True,
+    )
+    parser.add_argument(
+        "--num_padding_at_beginning",
+        type=int,
+        default=1,
+        help=
+        "OPT model has a fixed number (1) of padding tokens at the beginning of the input. We did not see this in other models but keep it as an option for now."
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--max_seq_len",
+        type=int,
+        default=512,
+        help="The maximum sequence length.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help=
+        "Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay",
+                       type=float,
+                       default=0.,
+                       help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs",
+                       type=int,
+                       default=1,
+                       help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help=
+        "Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="cosine",
+        help="The scheduler type to use.",
+        choices=[
+            "linear", "cosine", "cosine_with_restarts", "polynomial",
+            "constant", "constant_with_warmup"
+        ],
+    )
+    parser.add_argument(
+        "--num_warmup_steps",
+        type=int,
+        default=0,
+        help="Number of steps for the warmup in the lr scheduler.")
+    parser.add_argument("--output_dir",
+                       type=str,
+                       default=None,
+                       help="Where to store the model.")
+    parser.add_argument("--seed",
+                       type=int,
+                       default=1234,
+                       help="A seed for reproducible training.")
+    parser.add_argument("--local_rank",
+                       type=int,
+                       default=-1,
+                       help="local_rank for distributed training on gpus")
+    parser.add_argument('--gradient_checkpointing',
+                       action='store_true',
+                       help='Enable HF gradient checkpointing for model.')
+    parser.add_argument('--disable_dropout',
+                       action='store_true',
+                       help='Disable the dropout of the model.')
+    # deepspeed features
+    parser.add_argument('--offload',
+                       action='store_true',
+                       help='Enable ZeRO Offload techniques.')
+    parser.add_argument('--dtype',
+                       type=str,
+                       default='fp16',
+                       choices=['fp16', 'bf16'],
+                       help='Training data type')
+    parser.add_argument(
+        '--zero_stage',
+        type=int,
+        default=0,
+        help='ZeRO optimization stage for Actor model (and clones).')
+    ## LoRA for efficient training setting
+    parser.add_argument("--lora_dim",
+                       type=int,
+                       default=0,
+                       help="If > 0, use LoRA for efficient training.")
+    parser.add_argument("--lora_module_name",
+                       type=str,
+                       default="decoder.layers.",
+                       help="The scope of LoRA.")
+    parser.add_argument('--only_optimize_lora',
+                       action='store_true',
+                       help='Only optimize the LoRA parameters.')
+    parser.add_argument(
+        "--lora_learning_rate",
+        type=float,
+        default=5e-4,
+        help=
+        "Initial LoRA learning rate (after the potential warmup period) to use."
+    )
+    ## Tensorboard logging
+    parser.add_argument('--enable_tensorboard',
+                       action='store_true',
+                       help='Enable tensorboard logging')
+    parser.add_argument('--tensorboard_path',
+                       type=str,
+                       default="step2_tensorboard")
+    ## Print loss
+    parser.add_argument('--print_loss',
+                       action='store_true',
+                       help='Prints loss at each step.')
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args()
+
+    return args
+
+
+def main():
+    args = parse_args()
+
+    if args.local_rank == -1:
+        device = torch.device(get_accelerator().device_name())
+    else:
+        get_accelerator().set_device(args.local_rank)
+        device = torch.device(get_accelerator().device_name(), args.local_rank)
+        # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
+        # torch.distributed.init_process_group(backend='nccl')
+        deepspeed.init_distributed()
+
+    args.global_rank = torch.distributed.get_rank()
+
+    ds_config = get_train_ds_config(offload=args.offload,
+                                   dtype=args.dtype,
+                                   stage=args.zero_stage,
+                                   enable_tensorboard=args.enable_tensorboard,
+                                   tb_path=args.tensorboard_path,
+                                   tb_name="step2_model")
+    ds_config['train_micro_batch_size_per_gpu'] = args.per_device_train_batch_size
+    ds_config['train_batch_size'] = args.per_device_train_batch_size * torch.distributed.get_world_size() * args.gradient_accumulation_steps
+
+    set_random_seed(args.seed)
+    torch.distributed.barrier()
+
+    tokenizer = load_hf_tokenizer(args.model_name_or_path, fast_tokenizer=True)
+    # critic_model本质上是reward_model的一个副本, 是同一个模型的参数初始化得到的
+    rm_model = create_critic_model(args.model_name_or_path,
+                                  tokenizer,
+                                  ds_config,
+                                  args.num_padding_at_beginning,
+                                  disable_dropout=args.disable_dropout)
+
+    if args.lora_dim > 0:
+        rm_model = convert_linear_layer_to_lora(rm_model,
+                                               args.lora_module_name,
+                                               args.lora_dim)
+        if args.only_optimize_lora:
+            rm_model = only_optimize_lora_parameters(rm_model)
+
+    rm_model = make_model_gradient_checkpointing_compatible(rm_model)
+
+    # 设置当前为第二阶段的训练, 即Reward Model训练阶段
+    train_phase = 2
+    train_dataset, eval_dataset = create_prompt_dataset(args.local_rank, args.data_path, args.data_split,
+                                                       args.data_output_path, train_phase, args.seed,
+                                                       tokenizer, args.max_seq_len)
+    # 创建DataLoader, 在代码文件utils/data/data_utils.py中有具体实现DataCollatorReward类
+    data_collator = DataCollatorReward()
+    if args.local_rank == -1:
+        train_sampler = RandomSampler(train_dataset)
+        eval_sampler = SequentialSampler(eval_dataset)
+    else:
+        train_sampler = DistributedSampler(train_dataset)
+        eval_sampler = DistributedSampler(eval_dataset)
+    # 封装训练集数据迭代器
+    train_dataloader = DataLoader(train_dataset,
+                                 collate_fn=data_collator,
+                                 sampler=train_sampler,
+                                 batch_size=args.per_device_train_batch_size)
+    # 封装验证集数据迭代器
+    eval_sampler = SequentialSampler(eval_dataset)
+    eval_dataloader = DataLoader(eval_dataset,
+                                collate_fn=data_collator,
+                                sampler=eval_sampler,
+                                batch_size=args.per_device_eval_batch_size)
+
+    # 在main函数内部定义了价值评估函数
+    def evaluation_reward(model, eval_dataloader):
+        # 将模型设置为评估模式
+        model.eval()
+        # 初始化若干统计值为0
+        correct_predictions = 0
+        total_predictions = 0
+        scores = 0
+        for step, batch in enumerate(eval_dataloader):
+            batch = to_device(batch, device)
+            # 数据流必须禁止梯度计算和反向传播
+            with torch.no_grad():
+                outputs = model(**batch)
+            '''
+            outputs: {
+                'loss': tensor(), 
+                'chosen_mean_scores': tensor(batch_size,), 
+                'rejected_mean_scores': tensor(batch_size,)
+            }
+            '''
+            # chosen.shape: (batch_size,), rejected.shape: (batch_size, )
+            chosen = outputs["chosen_mean_scores"]
+            rejected = outputs["rejected_mean_scores"]
+            # chosen分值大于rejected分值, 即为赋分正确, 本质上就是"response的排序正确"
+            correct_predictions += (chosen > rejected).sum()
+            total_predictions += chosen.shape[0]
+            # 累加每个step的平均chosen分值
+            scores += outputs["chosen_mean_scores"].mean().float()
+            if step == 99: # For faster evaluation and debugging
+                break
+        # 计算acc, 和当前step的平均分数值
+        acc = correct_predictions / total_predictions
+        scores = scores / (step + 1)
+        try:
+            # 对多进程结果进行求和平均
+            acc = get_all_reduce_mean(acc).item()
+            scores = get_all_reduce_mean(scores).item()
+        except:
+            pass
+        # 最终返回平均分数值, acc值
+        return scores, acc
+
+    # 分组参数优化, 一部分参数采用weight decay策略, 另一部分不采用.
+    optimizer_grouped_parameters = get_optimizer_grouped_parameters(rm_model,
+                                                                   args.weight_decay,
+                                                                   args.lora_learning_rate)
+    # 实例化优化器对象
+    AdamOptimizer = DeepSpeedCPUAdam if args.offload else FusedAdam
+    optimizer = AdamOptimizer(optimizer_grouped_parameters,
+                            lr=args.learning_rate,
+                            betas=(0.9, 0.95))
+
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    # 实例化调节器对象
+    lr_scheduler = get_scheduler(name=args.lr_scheduler_type,
+                                optimizer=optimizer,
+                                num_warmup_steps=args.num_warmup_steps,
+                                num_training_steps=args.num_train_epochs *
+                                num_update_steps_per_epoch)
+
+    # 利用deepspeed封装model, 优化器, 调节器和参数, 加速训练!
+    rm_model, optimizer, _, lr_scheduler = deepspeed.initialize(model=rm_model,
+                                                               optimizer=optimizer,
+                                                               args=args,
+                                                               config=ds_config,
+                                                               lr_scheduler=lr_scheduler,
+                                                               dist_init_required=True)
+
+    if args.gradient_checkpointing:
+        rm_model.gradient_checkpointing_enable()
+
+    # 开始训练!!!
+    print_rank_0("***** Running training *****", args.global_rank)
+
+    print_rank_0(f"***** Evaluating reward, Epoch {0}/{args.num_train_epochs} *****", args.global_rank)
+    # 评估reward_model的表现
+    reward_score, acc = evaluation_reward(rm_model, eval_dataloader)
+    print_rank_0(f"chosen_last_scores (higher is better) : {reward_score}, acc (higher is better) : {acc}", args.global_rank)
+
+    # 经典的双重for循环训练模式
+    for epoch in range(args.num_train_epochs):
+        print_rank_0(f"Beginning of Epoch {epoch+1}/{args.num_train_epochs}, Total Micro Batches {len(train_dataloader)}", args.global_rank)
+        # 设置reward model为训练模式
+        rm_model.train()
+        mean_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            batch = to_device(batch, device)
+            # reward model进行前向传播计算出损失值
+            outputs = rm_model(**batch, use_cache=False)
+            '''
+            outputs: {
+                'loss': tensor(),
+                'chosen_mean_scores': tensor(batch_size,),
+                'rejected_mean_scores': tensor(batch_size,)
+            }
+            '''
+            loss = outputs["loss"]
+            # 经典"老三样", reward model进行反向传播
+            rm_model.backward(loss)
+            rm_model.step()
+            # 累加损失值, 并打印信息
+            mean_loss += loss.item()
+        print_rank_0(f"Epoch {epoch+1}/{args.num_train_epochs} with loss {mean_loss/(step+1)}", args.global_rank)
+        print_rank_0(f"***** Evaluating reward, Epoch {epoch+1}/{args.num_train_epochs} *****", args.global_rank)
+        # 在验证集上进行reward model的评估
+        reward_score, acc = evaluation_reward(rm_model, eval_dataloader)
+        print_rank_0(f"chosen_last_scores (higher is better) : {reward_score}, acc (higher is better) : {acc}", args.global_rank)
+        rm_model.tput_timer.update_epoch_count()
+
+    if args.output_dir is not None:
+        print_rank_0('saving model ...', args.global_rank)
+        rm_model = convert_lora_to_linear_layer(rm_model)
+
+        if args.global_rank == 0:
+            # save_hf_format(rm_model, tokenizer, args)
+            # 因为Qwen3大模型是以safetensor格式保存的, 所以需要重写模型保存的代码
+            save_hf_format_safetensors(rm_model, tokenizer, args)
+        if args.zero_stage == 3:
+            save_zero_three_model(rm_model,
+                                args.global_rank,
+                                args.output_dir,
+                                zero_stage=args.zero_stage)
+
+if __name__ == '__main__':
+    main()

RM-EN-01-30-2026/code/model_utils.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import math
+import time
+import torch
+from transformers import (
+    AutoConfig,
+    AutoModel,
+)
+from huggingface_hub import snapshot_download
+from transformers.integrations import HfDeepSpeedConfig
+
+from .reward_model import RewardModel
+from ..utils import load_state_dict_into_model
+
+
+def configure_dropout(model_config, dropout):
+    if dropout is not None:
+        for key in ('dropout', 'attention_dropout', 'hidden_dropout',
+                    'activation_dropout'):
+            if hasattr(model_config, key):
+                print(f"Setting model_config.{key} to {dropout}")
+                setattr(model_config, key, dropout)
+
+
+def causal_lm_model_to_fp32_loss(model):
+    """ Convert CausalLM model to calculate loss in fp32 """
+
+    def causal_lm_forward(
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **deprecated_arguments,
+    ):
+        kwargs = dict() if model.config.model_type == "llama" else dict(
+            head_mask=head_mask)
+        output = model.__original_forward__(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            labels=None,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs)
+
+        return_dict = isinstance(output, dict)
+        lm_logits = output.logits if return_dict else output[0]
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(lm_logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].float().contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            batch_size, seq_length, vocab_size = shift_logits.shape
+            # Flatten the tokens
+            loss_fct = torch.nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(batch_size * seq_length, vocab_size),
+                shift_labels.view(batch_size * seq_length))
+
+        if not return_dict:
+            # re-pack output with fp32 loss
+            return ((loss, ) + output) if loss is not None else output
+
+        output.loss = loss
+        return output
+
+    model.__original_forward__ = model.forward
+    model.forward = causal_lm_forward
+
+
+def create_hf_model(model_class,
+                    model_name_or_path,
+                    tokenizer,
+                    ds_config=None,
+                    rlhf_training=False,
+                    dropout=None):
+    model_config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
+    configure_dropout(model_config, dropout)
+
+    # Note: dschf is defined in function scope to avoid global effects
+    # https://huggingface.co/docs/transformers/main_classes/deepspeed#nontrainer-deepspeed-integration
+    if ds_config is not None and ds_config["zero_optimization"]["stage"] == 3:
+        dschf = HfDeepSpeedConfig(ds_config)
+    else:
+        dschf = None
+    if rlhf_training:
+        # the weight loading is handled by create critic model
+        with no_init_weights():
+            model = model_class.from_config(model_config)
+    else:
+        from transformers import AutoModelForCausalLM as _AutoModel
+        model = _AutoModel.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=True,
+            torch_dtype="auto",
+            device_map=None)
+
+    model.config.end_token_id = tokenizer.eos_token_id
+    model.config.pad_token_id = model.config.eos_token_id
+    model.resize_token_embeddings(int(
+        8 *
+        math.ceil(len(tokenizer) / 8.0)))  # make the vocab size multiple of 8
+
+    return model
+
+def create_critic_model(model_name_or_path,
+                       tokenizer,
+                       ds_config,
+                       num_padding_at_beginning=0,
+                       rlhf_training=False,
+                       disable_dropout=False,
+                       zero_stage=0):
+    start = time.time()
+    # 创建critic_model, 本质上也是调用上面的create_hf_model()函数
+    # 使用 AutoModelForCausalLM 加载，然后提取 .model（基础 transformer）
+    from transformers import AutoModelForCausalLM
+    full_model = create_hf_model(AutoModelForCausalLM, model_name_or_path, tokenizer,
+                                  ds_config, rlhf_training, disable_dropout)
+    # 提取基础 transformer 部分（返回 hidden_states 而非 logits）
+    if hasattr(full_model, 'model'):
+        critic_model = full_model.model  # Qwen3, LLaMA 等
+    elif hasattr(full_model, 'transformer'):
+        critic_model = full_model.transformer  # GPT-2 等
+    else:
+        critic_model = full_model
+    end = time.time()
+    # 单独运行第二阶段训练Reward Model的评估代码run_eval.sh时, 可能有报错, 可以暂时先注释下面两行即可
+    if torch.distributed.get_rank() == 0:
+        print(f"> Creating model from_config took {end - start} seconds")
+    
+    critic_model = RewardModel(critic_model,
+                              tokenizer,
+                              num_padding_at_beginning=num_padding_at_beginning)
+    
+    if rlhf_training:
+        # load critic model from checkpoint
+        if not os.path.isdir(model_name_or_path):
+            model_name_or_path = snapshot_download(model_name_or_path)
+        model_ckpt_path = os.path.join(model_name_or_path, 'pytorch_model.bin')
+        assert os.path.exists(model_ckpt_path), f"Cannot find model checkpoint at {model_ckpt_path}"
+        
+        start = time.time()
+        model_ckpt_state_dict = torch.load(model_ckpt_path, map_location='cpu')
+        end = time.time()
+        # 单独运行第二阶段训练Reward Model的评估代码run_eval.sh时, 有报错, 可以暂时先注释下面两行即可
+        if torch.distributed.get_rank() == 0:
+            print(f"> torch.load took {end - start} seconds")
+        
+        # load critic model from checkpoint with zero-stage 3 compatibility
+        # this functionality may be moved to DS checkpoint load API in future
+        start = time.time()
+        load_state_dict_into_model(critic_model,
+                                  model_ckpt_state_dict,
+                                  "",
+                                  zero_stage=zero_stage)
+        end = time.time()
+        # 单独运行第二阶段训练Reward Model的评估代码run_eval.sh时, 有报错, 可以暂时先注释下面两行即可
+        if torch.distributed.get_rank() == 0:
+            print(f"> Loading model state dict took {end - start} seconds")
+    
+    return critic_model

RM-EN-01-30-2026/code/raw_datasets.py

@@ -0,0 +1,828 @@
+# Copyright (c) Microsoft Corporation.
+from datasets import DatasetDict
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+# DeepSpeed Team
+from datasets import load_dataset, load_from_disk
+from torch.utils.data import Subset
+import re
+
+
+# The template prompt dataset class that all new dataset porting needs to
+# follow in order to have a unified API and unified data format.
+class PromptRawDataset(object):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        self.output_path = output_path
+        self.seed = seed
+        self.local_rank = local_rank
+        #if os.path.exists(dataset_name):
+        #    self.raw_datasets = load_from_disk(dataset_name)
+        if not dataset_name == 'local/jsonfile':
+            #self.raw_datasets = load_dataset(dataset_name)
+            self.raw_datasets = None
+
+    def get_train_data(self):
+        return
+
+    def get_eval_data(self):
+        return
+
+    # The prompt should be in the format of: " Human: " + actual_prompt_sentence + " Assistant:"
+    def get_prompt(self, sample):
+        return
+
+    # The chosen response should be in the format of: " " + actual_response_sentence
+    def get_chosen(self, sample):
+        return
+
+    # The rejected response should be in the format of: " " + actual_response_sentence
+    # If the dataset does not have rejected response, return None
+    def get_rejected(self, sample):
+        return
+
+    def get_prompt_and_chosen(self, sample):
+        return
+
+    def get_prompt_and_rejected(self, sample):
+        return
+
+
+# English dataset
+class DahoasRmstaticDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Dahoas/rm-static"
+        self.dataset_name_clean = "Dahoas_rm_static"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["test"]
+
+    def get_prompt(self, sample):
+        return sample['prompt']
+
+    def get_chosen(self, sample):
+        return sample['chosen']
+
+    def get_rejected(self, sample):
+        return sample['rejected']
+
+    def get_prompt_and_chosen(self, sample):
+        return sample['prompt'] + sample['chosen']
+
+    def get_prompt_and_rejected(self, sample):
+        return sample['prompt'] + sample['rejected']
+
+
+# English dataset
+class DahoasFullhhrlhfDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Dahoas/full-hh-rlhf"
+        self.dataset_name_clean = "Dahoas_full_hh_rlhf"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["test"]
+
+    def get_prompt(self, sample):
+        return sample['prompt']
+
+    def get_chosen(self, sample):
+        return sample['chosen']
+
+    def get_rejected(self, sample):
+        return sample['rejected']
+
+    def get_prompt_and_chosen(self, sample):
+        return sample['prompt'] + sample['chosen']
+
+    def get_prompt_and_rejected(self, sample):
+        return sample['prompt'] + sample['rejected']
+
+
+# English dataset
+class DahoasSyntheticinstructgptjpairwiseDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Dahoas/synthetic-instruct-gptj-pairwise"
+        self.dataset_name_clean = "Dahoas_synthetic_instruct_gptj_pairwise"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['prompt'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['chosen']
+
+    def get_rejected(self, sample):
+        return " " + sample['rejected']
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['prompt'] + " Assistant: " + sample['chosen']
+
+    def get_prompt_and_rejected(self, sample):
+        return " Human: " + sample['prompt'] + " Assistant: " + sample[
+            'rejected']
+
+
+# English dataset
+class YitingxieRlhfrewarddatasetsDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "yitingxie/rlhf-reward-datasets"
+        self.dataset_name_clean = "yitingxie_rlhf_reward_datasets"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["test"]
+
+    def get_prompt(self, sample):
+        return sample['prompt'] + "Assistant:"
+
+    def get_chosen(self, sample):
+        return sample['chosen'].split("Assistant:")[-1]
+
+    def get_rejected(self, sample):
+        return sample['rejected'].split("Assistant:")[-1]
+
+    def get_prompt_and_chosen(self, sample):
+        return sample['prompt'] + sample['chosen']
+
+    def get_prompt_and_rejected(self, sample):
+        return sample['prompt'] + sample['rejected']
+
+
+# English dataset
+class OpenaiWebgptcomparisonsDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "openai/webgpt_comparisons"
+        self.dataset_name_clean = "openai_webgpt_comparisons"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['question']['full_text'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        if float(sample['score_0']) >= float(sample['score_1']):
+            response = sample['answer_0']
+        else:
+            response = sample['answer_1']
+        # This data has citation square brackets and numbers (e.g., "[1]").
+        # Right now we are not doing browser-assisted finetuning, thus we
+        # remove these citations to avoid confusing the model.
+        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
+        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
+        return " " + response
+
+    def get_rejected(self, sample):
+        if float(sample['score_0']) < float(sample['score_1']):
+            response = sample['answer_0']
+        else:
+            response = sample['answer_1']
+        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
+        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
+        return " " + response
+
+    def get_prompt_and_chosen(self, sample):
+        if float(sample['score_0']) >= float(sample['score_1']):
+            response = sample['answer_0']
+        else:
+            response = sample['answer_1']
+        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
+        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
+        return " Human: " + sample['question'][
+            'full_text'] + " Assistant: " + response
+
+    def get_prompt_and_rejected(self, sample):
+        if float(sample['score_0']) < float(sample['score_1']):
+            response = sample['answer_0']
+        else:
+            response = sample['answer_1']
+        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
+        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
+        return " Human: " + sample['question'][
+            'full_text'] + " Assistant: " + response
+
+
+# English dataset
+class StanfordnlpSHPDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "stanfordnlp/SHP"
+        self.dataset_name_clean = "stanfordnlp_SHP"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["validation"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['history'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        if int(sample["labels"]) == 1:
+            response = sample["human_ref_A"]
+        else:
+            response = sample["human_ref_B"]
+        return " " + response
+
+    def get_rejected(self, sample):
+        if int(sample["labels"]) == 1:
+            response = sample["human_ref_B"]
+        else:
+            response = sample["human_ref_A"]
+        return " " + response
+
+    def get_prompt_and_chosen(self, sample):
+        if int(sample["labels"]) == 1:
+            response = sample["human_ref_A"]
+        else:
+            response = sample["human_ref_B"]
+        return " Human: " + sample['history'] + " Assistant: " + response
+
+    def get_prompt_and_rejected(self, sample):
+        if int(sample["labels"]) == 1:
+            response = sample["human_ref_B"]
+        else:
+            response = sample["human_ref_A"]
+        return " Human: " + sample['history'] + " Assistant: " + response
+
+
+# English dataset
+class PvduySharegptalpacaoavicunaformatDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "pvduy/sharegpt_alpaca_oa_vicuna_format"
+        self.dataset_name_clean = "pvduy_sharegpt_alpaca_oa_vicuna_format"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["test"]
+
+    def get_prompt(self, sample):
+        if sample['prompt'] is not None and len(sample['prompt']) > 0:
+            return sample['prompt'].replace("USER", "Human").replace(
+                "ASSISTANT", "Assistant")
+        return None
+
+    def get_chosen(self, sample):
+        if sample['label'] is not None and len(sample['label']) > 0:
+            return " " + sample['label']
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['prompt'] is not None and sample['label'] is not None and len(
+                sample['prompt']) > 0 and len(sample['label']) > 0:
+            return sample['prompt'].replace("USER", "Human").replace(
+                "ASSISTANT", "Assistant") + " " + sample['label']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+class LocalJsonFileDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name, chat_path):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "local/jsonfile"
+        self.dataset_name_clean = "jsonfile"
+        self.raw_datasets = load_dataset('json',
+                                         data_files={
+                                             "train":
+                                             chat_path + '/data/train.json',
+                                             "eval":
+                                             chat_path + '/data/eval.json'
+                                         })
+
+    def get_train_data(self):
+        if self.raw_datasets['train'] is not None:
+            return self.raw_datasets['train']
+        return None
+
+    def get_eval_data(self):
+        if self.raw_datasets['eval'] is not None:
+            return self.raw_datasets['eval']
+        return None
+
+    # The prompt should be in the format of: " Human: " + actual_prompt_sentence + " Assistant:"
+    def get_prompt(self, sample):
+        if sample['prompt'] is not None:
+            return " " + sample['prompt']
+        return None
+
+    # The chosen response should be in the format of: " " + actual_response_sentence
+    def get_chosen(self, sample):
+        if sample['chosen'] is not None:
+            return " " + sample['chosen']
+        return None
+
+    # The rejected response should be in the format of: " " + actual_response_sentence
+    # If the dataset does not have rejected response, return None
+    def get_rejected(self, sample):
+        if sample['rejected'] is not None:
+            return " " + sample['rejected']
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['prompt'] is not None and sample['chosen'] is not None:
+            return " " + sample['prompt'] + " " + sample['chosen']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        if sample['prompt'] is not None and sample['rejected'] is not None:
+            return " " + sample['prompt'] + " " + sample['rejected']
+        return None
+
+
+# Chinese dataset
+class Wangrui6ZhihuKOLDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "wangrui6/Zhihu-KOL"
+        self.dataset_name_clean = "wangrui6_Zhihu_KOL"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        if sample['INSTRUCTION'] is not None:
+            return " Human: " + sample['INSTRUCTION'] + " Assistant:"
+        return None
+
+    def get_chosen(self, sample):
+        if sample['RESPONSE'] is not None:
+            return " " + sample['RESPONSE']
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['INSTRUCTION'] is not None and sample['RESPONSE'] is not None:
+            return " Human: " + sample[
+                'INSTRUCTION'] + " Assistant: " + sample['RESPONSE']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Chinese dataset
+class CohereMiraclzhqueries2212Dataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Cohere/miracl-zh-queries-22-12"
+        self.dataset_name_clean = "Cohere_miracl_zh_queries_22_12"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["dev"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['query'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['positive_passages'][0]['text']
+
+    def get_rejected(self, sample):
+        return " " + sample['negative_passages'][0]['text']
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['query'] + " Assistant: " + sample[
+            'positive_passages'][0]['text']
+
+    def get_prompt_and_rejected(self, sample):
+        return " Human: " + sample['query'] + " Assistant: " + sample[
+            'negative_passages'][0]['text']
+
+
+# Chinese dataset
+class HelloSimpleAIHC3ChineseDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Hello-SimpleAI/HC3-Chinese"
+        self.dataset_name_clean = "Hello_SimpleAI_HC3_Chinese"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        if sample['question'] is not None:
+            return " Human: " + sample['question'] + " Assistant:"
+        return None
+
+    def get_chosen(self, sample):
+        if sample['human_answers'][0] is not None:
+            return " " + sample['human_answers'][0]
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['question'] is not None and sample['human_answers'][
+                0] is not None:
+            return " Human: " + sample['question'] + " Assistant: " + sample[
+                'human_answers'][0]
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Chinese dataset
+class MkqaChineseDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "mkqa-Chinese"
+        self.dataset_name_clean = "mkqa"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        if sample['queries']['zh_cn'] is not None:
+            return " Human: " + sample['queries']['zh_cn'] + " Assistant:"
+        return None
+
+    def get_chosen(self, sample):
+        if sample['answers']['zh_cn'][0]['text'] is not None:
+            return " " + sample['answers']['zh_cn'][0]['text']
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['queries']['zh_cn'] is not None and sample['answers'][
+                'zh_cn'][0]['text'] is not None:
+            return " Human: " + sample['queries'][
+                'zh_cn'] + " Assistant: " + sample['answers']['zh_cn'][0][
+                    'text']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Japanese dataset
+class MkqaJapaneseDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "mkqa-Japanese"
+        self.dataset_name_clean = "mkqa"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        if sample['queries']['ja'] is not None:
+            return " Human: " + sample['queries']['ja'] + " Assistant:"
+        return None
+
+    def get_chosen(self, sample):
+        if sample['answers']['ja'][0]['text'] is not None:
+            return " " + sample['answers']['ja'][0]['text']
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['queries']['ja'] is not None and sample['answers']['ja'][0][
+                'text'] is not None:
+            return " Human: " + sample['queries'][
+                'ja'] + " Assistant: " + sample['answers']['ja'][0]['text']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Japanese dataset
+class CohereMiracljaqueries2212Dataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Cohere/miracl-ja-queries-22-12"
+        self.dataset_name_clean = "Cohere_miracl_ja_queries_22_12"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["dev"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['query'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['positive_passages'][0]['text']
+
+    def get_rejected(self, sample):
+        return " " + sample['negative_passages'][0]['text']
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['query'] + " Assistant: " + sample[
+            'positive_passages'][0]['text']
+
+    def get_prompt_and_rejected(self, sample):
+        if len(sample['negative_passages']) > 0:
+            return " Human: " + sample['query'] + " Assistant: " + sample[
+                'negative_passages'][0]['text']
+        return None
+
+
+# Japanese dataset
+class LmqgQgjaquadDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "lmqg/qg_jaquad"
+        self.dataset_name_clean = "lmqg_qg_jaquad"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["validation"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['question'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['sentence']
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['question'] + " Assistant: " + sample[
+            'sentence']
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Japanese dataset
+class LmqgQagjaquadDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "lmqg/qag_jaquad"
+        self.dataset_name_clean = "lmqg_qag_jaquad"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["validation"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['questions'][0] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['paragraph']
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['questions'][0] + " Assistant: " + sample[
+            'paragraph']
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+# CustomDataset: 自定义数据集类，用于训练个性化垂直领域大模型，继承基类PromptRawDataset
+class CustomDataset(PromptRawDataset):
+    def __init__(self, output_path, seed, local_rank, dataset_name, chat_path):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        # 个性化数据集的名字可以自定义
+        self.dataset_name = "custom"
+        self.dataset_name_clean = "custom"
+        # 设定要读取的数据集所在的绝对路径
+        train_path = chat_path + '/data/train.jsonl'
+        eval_path = chat_path + '/data/dev.jsonl'
+        # 通过DatasetDict的类封装数据, 和load_dataset()函数保持一致.
+        self.raw_datasets = DatasetDict.from_json({'train': train_path, 'eval': eval_path})
+    
+    # 返回训练集数据
+    def get_train_data(self):
+        if self.raw_datasets['train'] is not None:
+            return self.raw_datasets['train']
+        return None
+    
+    # 返回验证集数据
+    def get_eval_data(self):
+        if self.raw_datasets['eval'] is not None:
+            return self.raw_datasets['eval']
+        return None
+    
+    # 构造prompt输入模型的格式: Human: prompt Assistant:
+    def get_prompt(self, sample):
+        if sample['prompt'] is not None:
+            return " Human: " + sample['prompt'] + " Assistant:"
+        return None
+    
+    # 构造chosen输入模型的格式: chosen
+    def get_chosen(self, sample):
+        if sample['chosen'] is not None:
+            return " " + sample['chosen']
+        return None
+    
+    # 构造reject输入模型的格式: reject
+    def get_rejected(self, sample):
+        if sample['rejected'] is not None:
+            return " " + sample['rejected']
+        return None
+    
+    # 构造第二阶���训练Reward Model的输入模型格式: Human: prompt Assistant: chosen
+    def get_prompt_and_chosen(self, sample):
+        if sample['prompt'] is not None and sample['chosen'] is not None:
+            return " Human: " + sample['prompt'] + " Assistant: " + sample['chosen']
+        return None
+    
+    # 构造第二阶段训练Reward Model的输入模型格式: Human: prompt Assistant: reject
+    def get_prompt_and_rejected(self, sample):
+        if sample['prompt'] is not None and sample['rejected'] is not None:
+            return " Human: " + sample['prompt'] + " Assistant: " + sample['rejected']
+        return None

RM-EN-01-30-2026/code/reward_model.py

@@ -0,0 +1,204 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+import torch
+from torch import nn
+
+# RewardModel这个类既实现了训练时用的forward() 的方法, 也实现了推理时用的forward_value()
+class RewardModel(nn.Module):
+    def __init__(self, base_model, tokenizer, num_padding_at_beginning=0):
+        super().__init__()
+        self.config = base_model.config
+        self.num_padding_at_beginning = num_padding_at_beginning
+        if hasattr(self.config, "word_embed_proj_dim"):
+            # OPT models use word_embed_proj_dim as final output
+            # 设置v_head将基于主干网络的输出特征hidden_size进行分值预测, 共输出max_seq_len个分数
+            self.v_head = nn.Linear(self.config.word_embed_proj_dim,
+                                   1,
+                                   bias=False)
+        else:
+            self.config.n_embd = self.config.hidden_size if hasattr(self.config, "hidden_size") else self.config.n_embd
+            # 同上, 设置v_head将基于主干网络的输出特征hidden_size进行分值预测, 共输出max_seq_len个分数
+            self.v_head = nn.Linear(self.config.n_embd, 1, bias=False)
+        # base_model即主干网络, 因此RM最终由1个主干网络 + 1个线性层v_head构成
+        self.rwtranrsformer = base_model
+        self.PAD_ID = tokenizer.pad_token_id
+        self.compute_fp32_loss = False
+
+    def gradient_checkpointing_enable(self):
+        self.rwtranrsformer.gradient_checkpointing_enable()
+
+    def gradient_checkpointing_disable(self):
+        self.rwtranrsformer.gradient_checkpointing_disable()
+
+    # 第2阶段调用的forward函数, 用于获取环境奖励与排序损失值 (区别于下面的forward_value()函数, 用于第3阶段推理调用)
+    def forward(self,
+                input_ids=None,
+                past_key_values=None,
+                attention_mask=None,
+                position_ids=None,
+                head_mask=None,
+                inputs_embeds=None,
+                use_cache=False):
+        loss = None
+        if self.config.model_type == "llama":
+            kwargs = dict()
+        else:
+            kwargs = dict(head_mask=head_mask)
+        # 此处的rwtransformer即为base_model基座模型, 也可以理解为主干网络
+        transformer_outputs = self.rwtranrsformer(input_ids,
+                                                 past_key_values=past_key_values,
+                                                 attention_mask=attention_mask,
+                                                 inputs_embeds=inputs_embeds,
+                                                 use_cache=use_cache,
+                                                 **kwargs)
+        # base_model模型输出一个列表, 位置0存储最后一层的输出张量
+        # hidden_states.shape: (batch_size * 2, max_seq_len, hidden_size)
+        hidden_states = transformer_outputs[0]
+        # 通过v_head线性层映射, 将最后一个维度从hidden_size降维成1, 并直接squeeze去掉
+        # rewards.shape: (batch_size * 2, max_seq_len), 相当于为序列中每一个位置的token都预测了一个价值
+        rewards = self.v_head(hidden_states).squeeze(-1)
+        chosen_mean_scores = []
+        rejected_mean_scores = []
+        # 确认input_ids是一个二维张量
+        assert len(input_ids.shape) == 2
+        # 在data_utils.py代码中, DataCollatorReward类中, 一个batch_size的数据被组装成了两部分, 实际的batch_size大小应该是输入的一半
+        # 此处要将其一分为二, 切分成chosen部分和reject部分
+        bs = input_ids.shape[0] // 2
+        seq_len = input_ids.shape[1]
+        # 切分出前半部分的chosen, 和后半部分的rejected
+        # 4个张量的shape均为: (batch_size, max_seq_len)
+        chosen_ids = input_ids[:bs]
+        rejected_ids = input_ids[bs:]
+        chosen_rewards = rewards[:bs]
+        rejected_rewards = rewards[bs:]
+        
+        # 计算Pairwise Ranking Loss
+        loss = 0
+        for i in range(bs):
+            # 取出同组chosen和rejected的token_id和分值reward
+            # chosen_id.shape: (max_seq_len, )
+            chosen_id = chosen_ids[i]
+            rejected_id = rejected_ids[i]
+            chosen_reward = chosen_rewards[i]
+            rejected_reward = rejected_rewards[i]
+            
+            # 下面的代码虽然看起来复杂, 但实质上在计算一个分割点
+            # c_ind为chosen_sentence的answer后的第一个pad_token的index
+            # 例如pad_token_id = 0, sentence = [1, 2, 3, 4, 5, 6, 0, 0, 0, 0]
+            # c_ind即为第一个pad_token的index = 6 
+            c_inds = (chosen_id == self.PAD_ID).nonzero()
+            c_ind = c_inds[self.num_padding_at_beginning].item() if len(c_inds) > self.num_padding_at_beginning else seq_len
+            
+            check_divergence = (chosen_id != rejected_id).nonzero()
+            # divergence_ind: 取chosen和rejected第一个不同的地方的index
+            # 可以理解为: response 中两个回答自由发挥的第1个token的index
+            if len(check_divergence) == 0:
+                end_ind = rejected_reward.size(-1)
+                divergence_ind = end_ind - 1
+                r_ind = c_ind
+            else:
+                # r_ind同理, 为reject_sentence的answer后的第一个pad_token的index
+                r_inds = (rejected_id == self.PAD_ID).nonzero()
+                r_ind = r_inds[self.num_padding_at_beginning].item() if len(r_inds) > self.num_padding_at_beginning else seq_len
+                # 较大者作为end_ind
+                end_ind = max(c_ind, r_ind)
+                divergence_ind = check_divergence[0]
+            assert divergence_ind > 0
+            
+            # AI图灵君课堂 (小朱老师独家讲义)
+            # 以chosen_sentence和reject_sentence最先不同的地方为起始, 生成结束的地方为终止
+            # 取两者在这个片段的对应分值, 这部分其实就是上个代码块提及的"对齐部分"
+            '''
+            max_seq_len为10, pad_token_id为0, 有同属同个prompt的chosen_sentence和reject_sentence:
+            prompt: [1, 2, 3]
+            chosen_sentence: [1, 2, 3, 4, 5, 6, 0, 0, 0, 0]
+            reject_sentence: [1, 2, 3, 7, 8, 0, 0, 0, 0, 0]
+            "两者answer的对齐部分", 即为"非prompt部分, 也非padding部分, 但长度要对齐":
+            chosen_truncated: [4, 5, 6]
+            reject_truncated: [7, 8, 0] 
+            '''
+            c_truncated_reward = chosen_reward[divergence_ind:end_ind]
+            r_truncated_reward = rejected_reward[divergence_ind:end_ind]
+            
+            # 下面的loss计算采用了"整个对齐部分的reward"来计算成对排序损失, 但是代码中对一个对话的预测评分实际上取的是该对话文本最后一个有效token的reward,
+            # 这个DeepSpeed团队也在论文中给出了说明, 这是一个开放性的策略, 用户可以自己制定个性化的评分策略, 比如answer部分的平均reward, 序列reward再接全连接层进行聚合后的reward, 等等
+            # 取代表结束的pad token所在位置的前一个位置(最后一个有效token的位置)的分值作为参考分值
+            chosen_mean_scores.append(chosen_reward[c_ind - 1])
+            rejected_mean_scores.append(rejected_reward[r_ind - 1])
+            
+            # 核心代码: 计算损失时使用了rank loss的形式, 是对chosen和rejected"对齐片段"进行计算的
+            # 计算采用了原始论文中的公式, 先计算sigmoid, 再进行log计算, 最终利用平均值作为损失值
+            # (c_truncated_reward - r_truncated_reward).shape: (truncated_seq_len,)
+            loss += -torch.nn.functional.logsigmoid(c_truncated_reward - r_truncated_reward).mean()
+        
+        loss = loss / bs
+        # 将batch_size个对话的reward值进行stack堆叠, chosen_mean_scores.shape: (batch_size, )
+        chosen_mean_scores = torch.stack(chosen_mean_scores)
+        rejected_mean_scores = torch.stack(rejected_mean_scores)
+        # 模型的返回字典中包含3个字段, loss, chosen分值, rejected分值
+        return {"loss": loss,
+                "chosen_mean_scores": chosen_mean_scores,
+                "rejected_mean_scores": rejected_mean_scores}
+    
+    # 第3阶段调用的推理函数-forward_value函数, 用于取到环境奖励和价值估计的方法
+    def forward_value(self,
+                     input_ids=None,
+                     attention_mask=None,
+                     past_key_values=None,
+                     position_ids=None,
+                     head_mask=None,
+                     inputs_embeds=None,
+                     return_value_only=False,
+                     prompt_length=0,
+                     use_cache=False):
+        '''
+        与forward的差别在于: forward需要针对输入的chosen-rejected对计算排序损失并返回,
+        而forward_value只需要考虑一个输入, 然后返回分值.
+        说白了, forward的输入是数据对, 因为要计算数据对的排序损失,
+        而forward_value的输入是单个数据, 直接推理出其分值.
+        return_value_only: 如果设置为True, 则在计算出values(在序列中每个token的分值预测)后直接返回.
+        '''
+        if self.config.model_type == "llama":
+            kwargs = dict()
+        else:
+            kwargs = dict(head_mask=head_mask)
+        # rwtransformer即base_model, 基座模型
+        transformer_outputs = self.rwtranrsformer(input_ids,
+                                                 past_key_values=past_key_values,
+                                                 attention_mask=attention_mask,
+                                                 inputs_embeds=inputs_embeds,
+                                                 use_cache=use_cache,
+                                                 **kwargs)
+        # [0]位置的张���即为base_model最后一层的输出张量
+        hidden_states = transformer_outputs[0]
+        # hidden_states.shape: (batch_size, max_seq_len, hidden_size)
+        # 经过线性层的映射, 在最后一个维度上, 每一个位置预测出一个分值
+        values = self.v_head(hidden_states).squeeze(-1)
+        # values.shape: (batch_size, max_seq_len)
+        
+        if return_value_only:
+            return values
+        else:
+            # [0 0 0 0 prompt, answer, 0 0 0 0 ] for step 3, we have padding at the beginning
+            # [prompt, answer, 0, 0, 0, 0] this is normal
+            assert prompt_length > 1, "prompt_length must be greater than 1 to help select the end score"
+            bs = values.size(0)
+            seq_len = input_ids.shape[1]
+            # 此变量的名称和作用, 与上面forward()函数中一致
+            chosen_end_scores = []
+            for i in range(bs):
+                input_id = input_ids[i]
+                value = values[i]
+                # value.shape: (max_seq_len)
+                # c_ind即为prompt之后的序列片段中, 第一个pad_token的index
+                c_inds = (input_id[prompt_length:] == self.PAD_ID).nonzero()
+                c_ind = c_inds[0].item() + prompt_length if len(c_inds) > 0 else seq_len
+                # 取c_ind的前一个index(实际上就是answer的最终位置)作为reward_score
+                chosen_end_scores.append(value[c_ind - 1])
+            # for循环结束后, len(chosen_end_scores) = batch_size, 相当于一个batch的样本分值
+            return {
+                "values": values,
+                "chosen_end_scores": torch.stack(chosen_end_scores)  # 经过stack堆叠后(batch_size,)
+            }

RM-EN-01-30-2026/data/rm_train.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20b4085690573224ca426fee9fc34363bb784b1bf46cf034016d17bd14b58c3a
+size 43901233

RM-EN-01-30-2026/model/chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

RM-EN-01-30-2026/model/config.json

@@ -0,0 +1,73 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "end_token_id": 151645,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9728,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 36,
+  "model_type": "qwen3",
+  "n_embd": 2560,
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "pad_token_id": 151645,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.0.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151672
+}

RM-EN-01-30-2026/model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1d6ad694f70c04ed664794dbf658e4c5d5f494efa3b2f1db1f400a316f4cf4e
+size 8043639192

RM-EN-01-30-2026/model/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650

RM-EN-01-30-2026/model/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "fast_tokenizer": true,
+  "is_local": true,
+  "model_max_length": 131072,
+  "pad_token": "<|im_end|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

RM-EN-01-30-2026/scripts/run_qwen3-4b.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+OUTPUT_DIR=./output_rm_en
+mkdir -p $OUTPUT_DIR
+
+deepspeed --num_gpus 1 main.py \
+    --model_name_or_path /workspace/Qwen3-4B \
+    --data_path custom \
+    --num_padding_at_beginning 0 \
+    --per_device_train_batch_size 2 \
+    --per_device_eval_batch_size 2 \
+    --max_seq_len 512 \
+    --learning_rate 1e-5 \
+    --weight_decay 0.1 \
+    --num_train_epochs 1 \
+    --gradient_accumulation_steps 8 \
+    --lr_scheduler_type cosine \
+    --num_warmup_steps 50 \
+    --seed 1234 \
+    --gradient_checkpointing \
+    --zero_stage 2 \
+    --offload \
+    --dtype bf16 \
+    --enable_tensorboard \
+    --tensorboard_path $OUTPUT_DIR/tensorboard \
+    --output_dir $OUTPUT_DIR \
+    --print_loss \
+    --deepspeed 2>&1 | tee $OUTPUT_DIR/training.log

SFT-EN-01-29-2026/README.md

@@ -0,0 +1,25 @@
+# SFT English Medical Model - Qwen3-4B
+
+## Overview
+- Base Model: Qwen3-4B
+- Training: DeepSpeed-Chat SFT with LoRA
+- Dataset: UltraMedical English (9K train, 1K eval)
+- Date: 2026-01-29
+
+## Training Config
+- LoRA dim: 64
+- Learning rate: 2e-5
+- Batch size: 2
+- Gradient accumulation: 4
+- ZeRO stage: 2
+- Dtype: bf16
+
+## Results
+- Final PPL: 2.498
+- Final Loss: 0.915
+
+## Directory
+- model/ - SFT model weights
+- data/ - Training data
+- scripts/ - Training scripts
+- code/ - Modified DeepSpeed-Chat code

SFT-EN-01-29-2026/code/data_utils.py

@@ -0,0 +1,629 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Part of the code was adopted from https://github.com/deepspeedai/Megatron-DeepSpeed/blob/main/megatron/data/dataset_utils.py
+"""
+import torch
+from torch.utils.data import Dataset, Subset, ConcatDataset
+from torch.nn.utils.rnn import pad_sequence
+import torch.nn.functional as F
+from datasets import load_dataset
+import numpy as np
+import os
+import hashlib
+from itertools import chain
+from dschat.utils.data import raw_datasets
+from deepspeed.accelerator import get_accelerator
+
+
+def get_raw_dataset(dataset_name, output_path, seed, local_rank):
+
+    if "Dahoas/rm-static" in dataset_name:
+        return raw_datasets.DahoasRmstaticDataset(output_path, seed,
+                                                  local_rank, dataset_name)
+    elif "Dahoas/full-hh-rlhf" in dataset_name:
+        return raw_datasets.DahoasFullhhrlhfDataset(output_path, seed,
+                                                    local_rank, dataset_name)
+    elif "Dahoas/synthetic-instruct-gptj-pairwise" in dataset_name:
+        return raw_datasets.DahoasSyntheticinstructgptjpairwiseDataset(
+            output_path, seed, local_rank, dataset_name)
+    elif "yitingxie/rlhf-reward-datasets" in dataset_name:
+        return raw_datasets.YitingxieRlhfrewarddatasetsDataset(
+            output_path, seed, local_rank, dataset_name)
+    elif "openai/webgpt_comparisons" in dataset_name:
+        return raw_datasets.OpenaiWebgptcomparisonsDataset(
+            output_path, seed, local_rank, dataset_name)
+    elif "stanfordnlp/SHP" in dataset_name:
+        return raw_datasets.StanfordnlpSHPDataset(output_path, seed,
+                                                  local_rank, dataset_name)
+    elif "pvduy/sharegpt_alpaca_oa_vicuna_format" in dataset_name:
+        return raw_datasets.PvduySharegptalpacaoavicunaformatDataset(
+            output_path, seed, local_rank, dataset_name)
+    elif "wangrui6/Zhihu-KOL" in dataset_name:
+        return raw_datasets.Wangrui6ZhihuKOLDataset(output_path, seed,
+                                                    local_rank, dataset_name)
+    elif "Cohere/miracl-zh-queries-22-12" in dataset_name:
+        return raw_datasets.CohereMiraclzhqueries2212Dataset(
+            output_path, seed, local_rank, dataset_name)
+    elif "Hello-SimpleAI/HC3-Chinese" in dataset_name:
+        return raw_datasets.HelloSimpleAIHC3ChineseDataset(
+            output_path, seed, local_rank, dataset_name)
+    elif "mkqa-Chinese" in dataset_name:
+        return raw_datasets.MkqaChineseDataset(output_path, seed, local_rank,
+                                               "mkqa")
+    elif "mkqa-Japanese" in dataset_name:
+        return raw_datasets.MkqaJapaneseDataset(output_path, seed, local_rank,
+                                                "mkqa")
+    elif "Cohere/miracl-ja-queries-22-12" in dataset_name:
+        return raw_datasets.CohereMiracljaqueries2212Dataset(
+            output_path, seed, local_rank, dataset_name)
+    elif "lmqg/qg_jaquad" in dataset_name:
+        return raw_datasets.LmqgQgjaquadDataset(output_path, seed, local_rank,
+                                                dataset_name)
+    elif "lmqg/qag_jaquad" in dataset_name:
+        return raw_datasets.LmqgQagjaquadDataset(output_path, seed, local_rank,
+                                                 dataset_name)
+    elif "local/jsonfile" in dataset_name:
+        chat_path = os.path.abspath(
+            os.path.join(os.path.dirname(__file__), os.path.pardir,
+                         os.path.pardir, os.path.pardir))
+        if not (os.path.isfile(chat_path + '/data/train.json')
+                and os.path.isfile(chat_path + '/data/eval.json')):
+            raise RuntimeError(
+                f"Please check both the train.json and eval.json files in your applications/DeepSpeed-Chat/data directory."
+            )
+        return raw_datasets.LocalJsonFileDataset(output_path, seed, local_rank,
+                                                 dataset_name, chat_path)
+    elif "custom" in dataset_name:
+        # 自動獲取當前文件所在的絕對路徑，向上跳三級到達 DeepSpeed-Chat 根目錄
+        current_file_path = os.path.dirname(os.path.abspath(__file__))
+        chat_path = os.path.abspath(os.path.join(current_file_path, os.path.pardir, os.path.pardir, os.path.pardir))
+        return raw_datasets.CustomDataset(output_path, seed, local_rank, 
+                                     dataset_name, chat_path)
+    else:
+        raise RuntimeError(
+            f"We do not have configs for dataset {dataset_name}, but you can add it by yourself in raw_datasets.py."
+        )
+
+
+def get_shuffle_idx(seed, size):
+    np_rng = np.random.RandomState(seed=seed)
+    dtype_ = np.uint32
+    if size >= (np.iinfo(np.uint32).max - 1):
+        dtype_ = np.int64
+    shuffle_idx = np.arange(start=0, stop=size, step=1, dtype=dtype_)
+    np_rng.shuffle(shuffle_idx)
+    return shuffle_idx
+
+# s=data_split, e.g., "6,2,2"
+def get_raw_dataset_split_index(local_rank,
+                                output_path,
+                                dataset_name,
+                                seed,
+                                split_name,
+                                data_split,
+                                split_index,
+                                data_size):
+    index_file_name = f"{output_path}/{dataset_name}_seed{seed}_{split_name}_{data_split}_{split_index}.npy"
+    # reindex each time when using local jsonfile since it's more likely to get modified
+    if (not os.path.isfile(index_file_name)) or (dataset_name
+                                                            == 'jsonfile'):
+        splits = [float(s) for s in data_split.split(',')]
+        splits_sum = sum(splits)
+        splits = [split / splits_sum for split in splits]
+        splits_index = [0]
+        for index, split in enumerate(splits):
+            splits_index.append(splits_index[index] +
+                                int(round(split * float(data_size))))
+        diff = splits_index[-1] - data_size
+        for index in range(1, len(splits_index)):
+            splits_index[index] -= diff
+        assert splits_index[-1] == data_size
+
+        shuffle_idx = get_shuffle_idx(seed, data_size)
+        for split_i in range(len(splits)):
+            shuffle_idx_split_file_name = f"{output_path}/{dataset_name}_seed{seed}_{split_name}_{data_split}_{split_i}.npy"
+            shuffle_idx_split = shuffle_idx[
+                splits_index[split_i]:splits_index[split_i + 1]]
+            np.save(shuffle_idx_split_file_name,
+                    shuffle_idx_split,
+                    allow_pickle=True)
+    index = np.load(index_file_name, allow_pickle=True)
+    return index.tolist()
+
+
+class PromptDataset(Dataset):
+
+    def __init__(self, prompt_dataset, chosen_dataset, reject_dataset,
+                 pad_token_id, train_phase) -> None:
+        super().__init__()
+        self.prompt_dataset = prompt_dataset
+        self.chosen_dataset = chosen_dataset
+        self.reject_dataset = reject_dataset
+        self.pad_token_id = pad_token_id
+        self.train_phase = train_phase
+
+    def __len__(self):
+        length = len(self.chosen_dataset)
+        if self.train_phase == 3:
+            length = len(self.prompt_dataset)
+        return length
+
+    def __getitem__(self, idx):
+        if self.train_phase == 1:
+            return {
+                "input_ids":
+                self.chosen_dataset[idx]["input_ids"],
+                "attention_mask":
+                self.chosen_dataset[idx]["attention_mask"],
+                "labels":self.chosen_dataset[idx]["input_ids"]
+                #torch.where(self.chosen_dataset[idx]["attention_mask"].bool(),
+                #            self.chosen_dataset[idx]["input_ids"], -100)
+            }
+        elif self.train_phase == 2:
+            return self.chosen_dataset[idx]["input_ids"], self.chosen_dataset[idx]["attention_mask"], \
+                self.reject_dataset[idx]["input_ids"], self.reject_dataset[idx]["attention_mask"]
+        elif self.train_phase == 3:
+            return self.prompt_dataset[idx]["input_ids"],self.prompt_dataset[idx]["attention_mask"], \
+                self.pad_token_id
+
+
+def create_dataset_split(current_dataset, raw_dataset, train_phase, tokenizer, 
+end_of_conversation_token, max_seq_len):
+    # 初始化3个空列表
+    prompt_dataset = []
+    chosen_dataset = []
+    reject_dataset = []
+    
+    # 如果当前为第一阶段训练, 即SFT阶段
+    if train_phase == 1:
+        # current_dataset: 传参进来的train_dataset, 已经经过Subset(train_dataset, 
+        # train_index)封装, 代表已经处理好的训练数据
+        for i, tmp_data in enumerate(current_dataset):
+            # 获取正常的(prompt, chosen)问答对, 用于第一阶段SFT训练
+            chosen_sentence = raw_dataset.get_prompt_and_chosen(tmp_data)
+            if chosen_sentence is not None:
+                # 对chosen_sentence尾部添加结束符
+                chosen_sentence += end_of_conversation_token
+                
+                # 对中文文本数据进行tokenizer处理, 本质就是text_to_id数字化的过程
+                chosen_token = tokenizer(chosen_sentence,
+                                        max_length=max_seq_len,
+                                        padding="max_length",
+                                        truncation=True,
+                                        return_tensors="pt")
+                
+                # 将input_ids和attention_mask字段取出, 并去掉batch_size=1的维度
+                chosen_token["input_ids"] = chosen_token["input_ids"].squeeze(0)
+                chosen_token["attention_mask"] = chosen_token["attention_mask"].squeeze(0)
+                chosen_dataset.append(chosen_token)
+    
+    
+    # 如果当前为第二阶段训练, 即Reward Model阶段
+    elif train_phase == 2:
+        for i, tmp_data in enumerate(current_dataset):
+            # 取出符合人类喜好的问答对(prompt, chosen)
+            chosen_sentence = raw_dataset.get_prompt_and_chosen(tmp_data)
+            # 取出不符合人类喜好的问答对(prompt, reject)
+            reject_sentence = raw_dataset.get_prompt_and_rejected(tmp_data)
+            
+            if chosen_sentence is not None and reject_sentence is not None:
+                # 在问答对的后面添加结束符
+                chosen_sentence += end_of_conversation_token
+                reject_sentence += end_of_conversation_token
+                
+                # 对符合人类喜好的问答对进行tokenizer处理, 并完成数字化id映射
+                chosen_token = tokenizer(chosen_sentence,
+                                        max_length=max_seq_len,
+                                        padding="max_length",
+                                        truncation=True,
+                                        return_tensors="pt")
+                
+                # 对不符合人类喜好的问答对进行tokenizer处理, 并完成数字化id映射
+                reject_token = tokenizer(reject_sentence,
+                                        max_length=max_seq_len,
+                                        padding="max_length",
+                                        truncation=True,
+                                        return_tensors="pt")
+                
+                # 将input_ids和attention_mask字段取出, 并添加进结果列表
+                chosen_token["input_ids"] = chosen_token["input_ids"]
+                chosen_token["attention_mask"] = chosen_token["attention_mask"]
+                chosen_dataset.append(chosen_token)
+                
+                reject_token["input_ids"] = reject_token["input_ids"]
+                reject_token["attention_mask"] = reject_token["attention_mask"]
+                reject_dataset.append(reject_token)
+    
+    # 如果当前为第三阶段训练, 即RLHF阶段
+    elif train_phase == 3:
+        # 不满足条件的数据, 直接过滤掉, 但需要统计被过滤掉的数据量
+        filtered = 0
+        for i, tmp_data in enumerate(current_dataset):
+            # 强化学习训练阶段, 只读取原始数据中的prompt输入
+            prompt = raw_dataset.get_prompt(tmp_data)
+            
+            if prompt is not None:
+                # 对prompt进行数字化映射和tokenizer处理
+                prompt_token = tokenizer(prompt, return_tensors="pt")
+                
+                # 只有数据长度满足条件的数据, 才需要被处理
+                # 如果length超过设定的最大序列长度(即max_prompt_len, 默认值256), 进行截断
+                if prompt_token["input_ids"].size()[-1] <= max_seq_len:
+                    for key_word in ["input_ids", "attention_mask"]:
+                        # 最后的 flip(0) 是将 token 序列进行 "翻转倒序"
+                        prompt_token[key_word] = prompt_token[key_word].squeeze(0).flip(0)
+                    
+                    # 一般来说, padding操作通常是直接在序列后面加入pad, padding后的输入序
+                    # 列变成了[prompt, padding]的形式, 那么自回归大模型将接在一连串pad后面继续生成, 这显然不合理.
+                    # 所以先将prompt进行flip(0)翻转倒序, 然后再padding, 达到符合条件的长
+                    # 度后最后再flip(0)翻转回来, 输入序列就变成了[padding, prompt]的形式, 大模型就可以接在prompt后面
+                    # 继续生成了.
+                    # 举个栗子: prompt_token_ids = [11, 22, 33], max_prompt_len = 5
+                    # 直接padding后, 就成了[11, 22, 33, 0, 0]
+                    # 如果先进行翻转倒序, prompt_token_ids.flip(0) = [33, 22, 11]
+                    # 再进行padding, prompt_token_ids.flip(0).padding() = [33, 22, 11, 0, 0]
+                    # 最后再次翻转倒序, prompt_token_ids.flip(0).padding().flip(0) = [0, 0, 11, 22, 33]
+                    '''
+                    注意: 最后一次翻转倒序是在data_utils.py代码文件中,
+                    DataCollatorRLHF()类内__call__函数内部进行的,
+                    batch["prompt"] = batch["prompt"].flip(1)
+                    因为此时已经是batch数据了, 所以翻转倒序是在flip(1), 即seq_len维度上进行的.
+                    '''
+                    prompt_dataset.append(prompt_token)
+                else:
+                    filtered += 1
+        
+        print(f'Creating dataset {raw_dataset.dataset_name_clean} '
+              f'for {train_phase=} size={len(prompt_dataset)} {filtered=}')
+    
+    # 返回封装类对象, 相当于torch中的Dataset, 供DataLoader调用
+    return PromptDataset(prompt_dataset, chosen_dataset, reject_dataset, 
+                        tokenizer.pad_token_id, train_phase)
+
+
+
+class PromptDataset(Dataset):
+    def __init__(self, prompt_dataset, chosen_dataset, reject_dataset,
+                 pad_token_id, train_phase) -> None:
+        super().__init__()
+        self.prompt_dataset = prompt_dataset
+        self.chosen_dataset = chosen_dataset
+        self.reject_dataset = reject_dataset
+        self.pad_token_id = pad_token_id
+        self.train_phase = train_phase
+    
+    def __len__(self):
+        length = len(self.chosen_dataset)
+        if self.train_phase == 3:
+            length = len(self.prompt_dataset)
+        return length
+    
+    def __getitem__(self, idx):
+        # 第一阶段SFT训练返回数据的格式
+        if self.train_phase == 1:
+            return {
+                "input_ids": self.chosen_dataset[idx]["input_ids"],
+                "attention_mask": self.chosen_dataset[idx]["attention_mask"],
+                "labels": self.chosen_dataset[idx]["input_ids"]
+            }
+        # 第二阶段Reward Model训练返回数据的格式
+        elif self.train_phase == 2:
+            return self.chosen_dataset[idx]["input_ids"], self.chosen_dataset[idx]["attention_mask"], \
+                   self.reject_dataset[idx]["input_ids"], self.reject_dataset[idx]["attention_mask"]
+        # 第三阶段RLHF训练返回数据的格式
+        elif self.train_phase == 3:
+            return self.prompt_dataset[idx]["input_ids"], self.prompt_dataset[idx]["attention_mask"], \
+                   self.pad_token_id
+
+
+def create_dataset(local_rank, dataset_name, data_split, output_path,
+                  train_phase, seed, tokenizer, end_of_conversation_token,
+                  max_seq_len):
+    # 训练个性化私有大模型, 设置dataset_name='custom'
+    dataset_name = "custom"
+    # 因为设定了dataset_name = 'custom', 所以调用get_raw_dataset()函数时, 就自动注册了
+    # custom分支, 从本地读取数据集
+    raw_dataset = get_raw_dataset(dataset_name, output_path, seed, local_rank)
+    
+    # 调用在CustomDataset类中定义的get_train_data()函数, 获取训练集数据
+    train_dataset = raw_dataset.get_train_data()
+    
+    # 获取随机排列下标后的训练集index列表对象
+    train_index = get_raw_dataset_split_index(local_rank, output_path,
+                                             raw_dataset.dataset_name_clean,
+                                             seed, "train", data_split,
+                                             train_phase - 1,
+                                             len(train_dataset))
+    
+    # 传参train_dataset数据集, 和随机排列后的train_index列表对象, 封装成Subset
+    # Subset功能: 取指定一个索引序列对应的子数据集
+    train_dataset = Subset(train_dataset, train_index)
+    
+    # 调用核心函数create_dataset_split()进行数据切分处理
+    train_dataset = create_dataset_split(train_dataset, raw_dataset,
+                                        train_phase, tokenizer,
+                                        end_of_conversation_token,
+                                        max_seq_len)
+    
+    # 下面验证集的数据处理流程, 同上面训练集一样
+    eval_dataset = raw_dataset.get_eval_data()
+    
+    eval_index = get_raw_dataset_split_index(local_rank, output_path,
+                                            raw_dataset.dataset_name_clean,
+                                            seed, "eval",
+                                            data_split, train_phase - 1,
+                                            len(eval_dataset))
+    
+    eval_dataset = Subset(eval_dataset, eval_index)
+    eval_dataset = create_dataset_split(eval_dataset, raw_dataset, train_phase,
+                                       tokenizer, end_of_conversation_token,
+                                       max_seq_len)
+    
+    return train_dataset, eval_dataset
+
+
+def create_prompt_dataset(local_rank,
+                          data_path,
+                          data_split,
+                          output_path,
+                          train_phase,
+                          seed,
+                          tokenizer,
+                          max_seq_len,
+                          end_of_conversation_token="<|endoftext|>",
+                          sft_only_data_path=[],
+                          reload=False):
+    """
+    Creates the prompt dataset
+    """
+    os.makedirs(output_path, exist_ok=True)
+    fname = "_".join(data_path)
+    sft_cache_key = "_".join(sft_only_data_path)
+    tokenizer_name = tokenizer.init_kwargs["name_or_path"].replace("/", "_")
+    fname = f"{fname}_split{data_split}_phase{train_phase}_seed{seed}_tokenizer{tokenizer_name}_seqlen{max_seq_len}_sft{sft_cache_key}"
+    fname = "_".join(fname.split("/"))
+    fname = hashlib.sha256(fname.encode()).hexdigest(
+    )  # hash the file name to avoid too long file name
+    train_fname = f"{output_path}/traindata_{fname}.pt"
+    eval_fname = f"{output_path}/evaldata_{fname}.pt"
+
+    cache_found = os.path.isfile(train_fname) and os.path.isfile(eval_fname)
+    buf_create_cache = torch.ByteTensor([not cache_found]).to(
+        get_accelerator().current_device_name())
+    torch.distributed.all_reduce(buf_create_cache)
+
+    if local_rank <= 0 and (buf_create_cache.item() != 0 or reload):
+        print(f'Creating prompt dataset {data_path}, {reload=}')
+        if len(data_path) == 1:  # Single dataset.
+            train_dataset, eval_dataset = create_dataset(
+                local_rank,
+                data_path[0],
+                data_split,
+                output_path,
+                train_phase,
+                seed,
+                tokenizer,
+                end_of_conversation_token,
+                max_seq_len,
+                )
+        else:  # Blending datasets.
+            train_datasets = []
+            eval_datasets = []
+            train_size = 0
+            eval_size = 0
+            for d_path in data_path:
+                train_dataset, eval_dataset = create_dataset(
+                    local_rank,
+                    d_path,
+                    data_split,
+                    output_path,
+                    train_phase,
+                    seed,
+                    tokenizer,
+                    end_of_conversation_token,
+                    max_seq_len,
+                    )
+                train_datasets.append(train_dataset)
+                eval_datasets.append(eval_dataset)
+                train_size += len(train_dataset)
+                eval_size += len(eval_dataset)
+            train_dataset = ConcatDataset(train_datasets)
+            shuffle_idx = get_shuffle_idx(seed, train_size)
+            train_dataset = Subset(train_dataset, shuffle_idx.tolist())
+            eval_dataset = ConcatDataset(eval_datasets)
+            shuffle_idx = get_shuffle_idx(seed, eval_size)
+            eval_dataset = Subset(eval_dataset, shuffle_idx.tolist())
+
+        # Append the SFT-only dataset if it exists, and current phase is 1(SFT).
+        if train_phase == 1 and sft_only_data_path:
+            sft_train_datasets = []
+            sft_eval_datasets = []
+            sft_train_size = 0
+            sft_eval_size = 0
+            for sft_path in sft_only_data_path:
+                sft_train_dataset, sft_eval_dataset = create_dataset(
+                    local_rank,
+                    sft_path,
+                    "10,0,0",
+                    output_path,
+                    train_phase,
+                    seed,
+                    tokenizer,
+                    end_of_conversation_token,
+                    max_seq_len,
+                    )
+                sft_train_datasets.append(sft_train_dataset)
+                sft_eval_datasets.append(sft_eval_dataset)
+                sft_train_size += len(sft_train_dataset)
+                sft_eval_size += len(sft_eval_dataset)
+            if sft_train_datasets:  # Check if sft_train_datasets is not empty
+                sft_train_dataset = ConcatDataset(sft_train_datasets)
+                train_dataset = ConcatDataset(
+                    [train_dataset, sft_train_dataset])
+                shuffle_idx = get_shuffle_idx(seed, len(train_dataset))
+                train_dataset = Subset(train_dataset, shuffle_idx.tolist())
+            if sft_eval_datasets:  # Check if sft_eval_datasets is not empty
+                sft_eval_dataset = ConcatDataset(sft_eval_datasets)
+                eval_dataset = ConcatDataset([eval_dataset, sft_eval_dataset])
+                shuffle_idx = get_shuffle_idx(seed, len(eval_dataset))
+                eval_dataset = Subset(eval_dataset, shuffle_idx.tolist())
+        torch.save(train_dataset, train_fname)
+        torch.save(eval_dataset, eval_fname)
+    torch.distributed.barrier()
+    return torch.load(train_fname,
+                      weights_only=False), torch.load(eval_fname,
+                                                      weights_only=False)
+
+
+class DataCollatorReward:
+
+    def __call__(self, data):
+        batch = {}
+        batch["input_ids"] = torch.cat([f[0]
+                                        for f in data] + [f[2] for f in data],
+                                       dim=0)
+        batch["attention_mask"] = torch.cat([f[1] for f in data] +
+                                            [f[3] for f in data],
+                                            dim=0)
+        return batch
+
+# 3. RLHF数据集的处理
+class DataCollatorRLHF:
+
+    def __init__(self, max_token_len, inference_tp_size):
+        self.max_token_len = max_token_len
+        self.inference_tp_size = inference_tp_size
+
+    def __call__(self, data):
+        batch = {}
+        pad_token_id = data[-1][-1]
+
+        prompt = pad_sequence([f[0] for f in data],
+                              padding_value=pad_token_id,
+                              batch_first=True)
+        prompt_mask = pad_sequence([f[1] for f in data],
+                                   padding_value=0,
+                                   batch_first=True)
+
+        ### make sure the final ouput is a seqence of 2**?
+        length = prompt.size()[-1]
+        pad_length = self.max_token_len - length
+        if pad_length > 0:
+            batch["prompt"] = F.pad(prompt,
+                                    pad=(0, pad_length),
+                                    mode='constant',
+                                    value=pad_token_id)
+            batch["prompt_att_mask"] = F.pad(prompt_mask,
+                                             pad=(0, pad_length),
+                                             mode='constant',
+                                             value=0)
+        else:
+            batch["prompt"] = prompt
+            batch["prompt_att_mask"] = prompt_mask
+        batch["prompt"] = batch["prompt"].flip(1)
+        batch["prompt_att_mask"] = batch["prompt_att_mask"].flip(1)
+        return batch
+
+
+def get_unsupervised_data(args, tokenizer):
+    unsupervised_raw_datasets = load_dataset(
+        args.unsupervised_dataset_name, args.unsupervised_dataset_config_name)
+    column_names = unsupervised_raw_datasets["train"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    def tokenize_function(examples):
+        return tokenizer(examples[text_column_name])
+
+    tokenized_datasets = unsupervised_raw_datasets.map(
+        tokenize_function,
+        batched=True,
+        num_proc=args.preprocessing_num_workers,
+        remove_columns=column_names,
+        load_from_cache_file=True,
+        desc="Running tokenizer on dataset",
+    )
+
+    block_size = args.max_prompt_seq_len + args.max_answer_seq_len
+
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {
+            k: list(chain(*examples[k]))
+            for k in examples.keys()
+        }
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= block_size:
+            total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        result = {
+            k:
+            [t[i:i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    lm_datasets = tokenized_datasets.map(
+        group_texts,
+        batched=True,
+        num_proc=args.preprocessing_num_workers,
+        load_from_cache_file=True,
+        desc=f"Grouping texts in chunks of {block_size}",
+    )
+
+    train_dataset = lm_datasets["train"]
+
+    return train_dataset
+
+
+class MiniDataset:
+
+    def __init__(self, max_size, small_batch_size):
+        self.dataset = []
+        self.max_size = max_size
+        self.small_batch_size = small_batch_size
+
+    def seperate(self):
+        small_dataset = []
+        for large_batch in self.dataset:
+            if type(large_batch) == list or type(large_batch) == tuple:
+                large_size = len(large_batch[0])
+            elif type(large_batch) == dict:
+                large_size = len(large_batch[list(large_batch.keys())[0]])
+            else:
+                large_size = len(large_batch)
+            for i in range(0, large_size, self.small_batch_size):
+                if type(large_batch) == list or type(large_batch) == tuple:
+                    small_dataset.append(
+                        [x[i:i + self.small_batch_size] for x in large_batch])
+                elif type(large_batch) == dict:
+                    small_dataset.append({
+                        k: v[i:i + self.small_batch_size]
+                        for k, v in large_batch.items()
+                    })
+                else:
+                    small_dataset.append(large_batch[i:i +
+                                                     self.small_batch_size])
+        self.free()
+
+        return small_dataset
+
+    def add(self, data):
+        if len(self.dataset) < self.max_size:
+            self.dataset.append(data)
+            if len(self.dataset) == self.max_size:
+                return self.seperate()
+            else:
+                return None
+        else:
+            raise ValueError(
+                "The dataset is full but we did not stop it. There is a bug in the code."
+            )
+
+    def free(self):
+        self.dataset = []

SFT-EN-01-29-2026/code/main.py

@@ -0,0 +1,866 @@
+#!/usr/bin/env python
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import argparse
+import math
+import sys
+sys.path.append("/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat")
+import torch
+from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
+from torch.utils.data.distributed import DistributedSampler
+
+from transformers import (
+    AutoModelForCausalLM,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+
+import deepspeed
+from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
+from deepspeed import get_accelerator
+
+from dschat.utils.data.data_utils import create_prompt_dataset
+from dschat.utils.utils import print_rank_0, to_device, save_hf_format, set_random_seed, get_all_reduce_mean, get_optimizer_grouped_parameters, save_zero_three_model, load_hf_tokenizer, save_hf_format_safetensors
+from dschat.utils.ds_utils import get_train_ds_config
+from dschat.utils.module.lora import convert_linear_layer_to_lora, convert_lora_to_linear_layer, only_optimize_lora_parameters, make_model_gradient_checkpointing_compatible
+from dschat.utils.model.model_utils import create_hf_model, causal_lm_model_to_fp32_loss
+from dschat.utils.perf import print_throughput
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description=
+        "Finetune a transformers model on a causal language modeling task")
+    # 默认的数据集我们不⽤, data_path传参进来的是垂直领域的个性化数据集
+    parser.add_argument('--data_path',
+                        nargs='*',
+                        default=['Dahoas/rm-static'],
+                        help='Path to the training dataset. Accepted format:'
+                        '1) a single data path, 2) multiple datasets in the'
+                        'form: dataset1-path dataset2-path ...')
+    parser.add_argument('--data_split',
+                        type=str,
+                        default='6,2,2',
+                        help='Comma-separated list of proportions for training'
+                        'phase 1, 2, and 3 data. For example the split `6,2,2`'
+                        'will use 60%% of data for phase 1, 20%% for phase 2'
+                        'and 20%% for phase 3.')
+    parser.add_argument(
+        '--sft_only_data_path',
+        nargs='*',
+        default=[],
+        help='Path to the dataset for only using in SFT phase.')
+    parser.add_argument(
+        '--data_output_path',
+        type=str,
+        default='/tmp/data_files/',
+        help=
+        'Where to store the data-related files such as shuffle index. This needs to be on a local storage of a node (not on a shared storage)'
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help=
+        "Path to pretrained model or model identifier from huggingface.co/models.",
+        required=True,
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--max_seq_len",
+        type=int,
+        default=512,
+        help="The maximum sequence length.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-3,
+        help=
+        "Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay",
+                        type=float,
+                        default=0.,
+                        help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs",
+                        type=int,
+                        default=1,
+                        help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help=
+        "Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="cosine",
+        help="The scheduler type to use.",
+        choices=[
+            "linear", "cosine", "cosine_with_restarts", "polynomial",
+            "constant", "constant_with_warmup"
+        ],
+    )
+    parser.add_argument(
+        "--num_warmup_steps",
+        type=int,
+        default=0,
+        help="Number of steps for the warmup in the lr scheduler.")
+    parser.add_argument("--output_dir",
+                        type=str,
+                        default=None,
+                        help="Where to store the model.")
+    parser.add_argument("--seed",
+                        type=int,
+                        default=1234,
+                        help="A seed for reproducible training.")
+    parser.add_argument("--local_rank",
+                        type=int,
+                        default=-1,
+                        help="local_rank for distributed training on gpus")
+    parser.add_argument('--gradient_checkpointing',
+                        action='store_true',
+                        help='Enable HF gradient checkpointing for model.')
+    parser.add_argument(
+        "--dropout",
+        type=float,
+        default=None,
+        help="If dropout configured, use it. "
+        "Otherwise, keep the default dropout configuration of the model.")
+    # deepspeed features
+    parser.add_argument('--offload',
+                        action='store_true',
+                        help='Enable ZeRO Offload techniques.')
+    parser.add_argument('--dtype',
+                        type=str,
+                        default='fp16',
+                        choices=['fp16', 'bf16'],
+                        help='Training data type')
+    parser.add_argument(
+        '--zero_stage',
+        type=int,
+        default=0,
+        help='ZeRO optimization stage for Actor model (and clones).')
+    ## LoRA for efficient training setting
+    parser.add_argument("--lora_dim",
+                        type=int,
+                        default=0,
+                        help="If > 0, use LoRA for efficient training.")
+    parser.add_argument("--lora_module_name",
+                        type=str,
+                        default="decoder.layers.",
+                        help="The scope of LoRA.")
+    parser.add_argument('--only_optimize_lora',
+                        action='store_true',
+                        help='Only optimize the LoRA parameters.')
+    parser.add_argument(
+        "--lora_learning_rate",
+        type=float,
+        default=5e-4,
+        help=
+        "Initial LoRA learning rate (after the potential warmup period) to use."
+    )
+    ## low precision
+    parser.add_argument(
+        '--compute_fp32_loss',
+        action='store_true',
+        help='Relevant for low precision dtypes (fp16, bf16, etc.). '
+        'If specified, loss is calculated in fp32.')
+    ## Tensorboard logging
+    parser.add_argument('--enable_tensorboard',
+                        action='store_true',
+                        help='Enable tensorboard logging')
+    parser.add_argument('--tensorboard_path',
+                        type=str,
+                        default="step1_tensorboard")
+    ## Tokenizer
+    parser.add_argument(
+        "--add_eot_token",
+        action='store_true',
+        help="Add `eot_token` as additional special token to tokenizer")
+    parser.add_argument(
+        "--eot_token",
+        type=str,
+        default="<|endoftext|>",
+        help="Specify the format of the `eot_token`",
+    )
+    ## Print loss
+    parser.add_argument('--print_loss',
+                        action='store_true',
+                        help='Prints loss at each step.')
+    # 此处是所有超参数和训练参数的设置位置
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args()
+
+    return args
+
+
+def main():
+    args = parse_args()
+
+    if args.local_rank == -1:
+        device = torch.device(get_accelerator().device_name())
+    else:
+        get_accelerator().set_device(args.local_rank)
+        device = torch.device(get_accelerator().device_name(), args.local_rank)
+        # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
+        # torch.distributed.init_process_group(backend='nccl')
+        deepspeed.init_distributed()
+
+    args.global_rank = torch.distributed.get_rank()
+
+    ds_config = get_train_ds_config(offload=args.offload,
+                                    dtype=args.dtype,
+                                    stage=args.zero_stage,
+                                    enable_tensorboard=args.enable_tensorboard,
+                                    tb_path=args.tensorboard_path,
+                                    tb_name="step1_model")
+    ds_config[
+        'train_micro_batch_size_per_gpu'] = args.per_device_train_batch_size
+    ds_config[
+        'train_batch_size'] = args.per_device_train_batch_size * torch.distributed.get_world_size(
+        ) * args.gradient_accumulation_steps
+
+    # If passed along, set the training seed now.
+    set_random_seed(args.seed)
+
+    torch.distributed.barrier()
+
+    # 实例化tokenizer和model
+    # load_hf_tokenizer will get the correct tokenizer and set padding tokens based on the model family
+    additional_special_tokens = args.eot_token if args.add_eot_token else None
+    tokenizer = load_hf_tokenizer(args.model_name_or_path,
+                                  fast_tokenizer=True,
+                                  add_special_tokens=additional_special_tokens)
+
+    model = create_hf_model(AutoModelForCausalLM,
+                            args.model_name_or_path,
+                            tokenizer,
+                            ds_config,
+                            dropout=args.dropout)
+
+    if args.compute_fp32_loss:
+        print_rank_0(
+            f"Using model {model.__class__.__name__} with loss in fp32",
+            args.global_rank)
+        causal_lm_model_to_fp32_loss(model)
+
+    # 设置LoRA微调
+    if args.lora_dim > 0:
+        model = convert_linear_layer_to_lora(model, args.lora_module_name,
+                                             args.lora_dim)
+        if args.only_optimize_lora:
+            model = only_optimize_lora_parameters(model)
+            model = make_model_gradient_checkpointing_compatible(model)
+
+    # 准备训练数据, 注意当前处于第⼀阶段 SFT
+    # Prepare the data
+    train_phase = 1
+    print('args: ', args)
+    print('data_path: ', args.data_path)
+    train_dataset, eval_dataset = create_prompt_dataset(
+        args.local_rank,
+        args.data_path,
+        args.data_split,
+        args.data_output_path,
+        train_phase,
+        args.seed,
+        tokenizer,
+        args.max_seq_len,
+        end_of_conversation_token=tokenizer.eos_token,
+        sft_only_data_path=args.sft_only_data_path)
+    # DataLoaders creation:
+    if args.local_rank == -1:
+        train_sampler = RandomSampler(train_dataset)
+        eval_sampler = SequentialSampler(eval_dataset)
+    else:
+        train_sampler = DistributedSampler(train_dataset)
+        eval_sampler = DistributedSampler(eval_dataset)
+    train_dataloader = DataLoader(train_dataset,
+                                  collate_fn=default_data_collator,
+                                  sampler=train_sampler,
+                                  batch_size=args.per_device_train_batch_size)
+    eval_dataloader = DataLoader(eval_dataset,
+                                 collate_fn=default_data_collator,
+                                 sampler=eval_sampler,
+                                 batch_size=args.per_device_eval_batch_size)
+
+    # main内部定义的评估函数
+    def evaluation(model, eval_dataloader):
+        model.eval()
+        losses = 0
+        for step, batch in enumerate(eval_dataloader):
+            batch = to_device(batch, device)
+            with torch.no_grad():
+                outputs = model(**batch)
+
+            loss = outputs.loss
+            losses += loss.float()
+        losses = losses / (step + 1)
+        try:
+            losses = get_all_reduce_mean(losses)
+        except:
+            pass
+        try:
+            perplexity = torch.exp(losses).item()
+        except OverflowError:
+            perplexity = float("inf")
+        return perplexity, losses.item()
+
+    # 采⽤分组优化参数的优化器策略
+    # Split weights in two groups, one with weight decay and the other not.
+    optimizer_grouped_parameters = get_optimizer_grouped_parameters(
+        model, args.weight_decay, args.lora_learning_rate)
+
+    AdamOptimizer = DeepSpeedCPUAdam if args.offload else FusedAdam
+    optimizer = AdamOptimizer(optimizer_grouped_parameters,
+                              lr=args.learning_rate,
+                              betas=(0.9, 0.95))
+
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader) / args.gradient_accumulation_steps)
+    lr_scheduler = get_scheduler(
+        name=args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=args.num_warmup_steps,
+        num_training_steps=args.num_train_epochs * num_update_steps_per_epoch,
+    )
+
+    # 采⽤deepspeed对相关组件进⾏封装, 本质上是进⾏加速优化
+    model, optimizer, _, lr_scheduler = deepspeed.initialize(
+        model=model,
+        optimizer=optimizer,
+        args=args,
+        config=ds_config,
+        lr_scheduler=lr_scheduler,
+        dist_init_required=True)
+
+    if args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+
+    # 开始训练, 打印⼀些关键信息
+    # Train!
+    print_rank_0("***** Running training *****", args.global_rank)
+#     print_rank_0(
+#         f"***** Evaluating perplexity, Epoch {0}/{args.num_train_epochs} *****",
+#         args.global_rank)
+#     perplexity, eval_loss = evaluation(model, eval_dataloader)
+#    print_rank_0(f"ppl: {perplexity}, loss: {eval_loss}", args.global_rank)
+
+    # 经典的双重for循环训练模式
+    for epoch in range(args.num_train_epochs):
+        print_rank_0(
+            f"Beginning of Epoch {epoch+1}/{args.num_train_epochs}, Total Micro Batches {len(train_dataloader)}",
+            args.global_rank)
+        # 将模型设置为训练模式
+        model.train()
+        import time
+        for step, batch in enumerate(train_dataloader):
+            start = time.time()
+            batch = to_device(batch, device)
+            # 模型的前向传播计算, 并取到损失值loss
+            outputs = model(**batch, use_cache=False)
+            loss = outputs.loss
+            if args.print_loss:
+                print(
+                    f"Epoch: {epoch}, Step: {step}, Rank: {torch.distributed.get_rank()}, loss = {loss}"
+                )
+            # 反向传播, "⽼三样"
+            model.backward(loss)
+            model.step()
+            end = time.time()
+            if torch.distributed.get_rank() == 0:
+                print_throughput(model.model, args, end - start,
+                                 args.global_rank)
+
+        # 在验证集上进⾏评估, 获取困惑度
+        # Evaluate perplexity on the validation set.
+        print_rank_0(
+            f"***** Evaluating perplexity, Epoch {epoch+1}/{args.num_train_epochs} *****",
+            args.global_rank)
+        perplexity, eval_loss = evaluation(model, eval_dataloader)
+        print_rank_0(f"ppl: {perplexity}, loss: {eval_loss}", args.global_rank)
+        model.tput_timer.update_epoch_count()
+
+    if args.output_dir is not None:
+        print_rank_0('saving the final model ...', args.global_rank)
+        model = convert_lora_to_linear_layer(model)
+
+        if args.global_rank == 0:
+            # save_hf_format(model, tokenizer, args)
+            # 因为我们项⽬中需要训练的是Qwen3⼤模型, 需要保存成safetensor的格式
+            save_hf_format_safetensors(model, tokenizer, args)
+
+        if args.zero_stage == 3:
+            # 在zero_stage==3时, 每⼀个GPU只包含model的⼀部分, 因此需要⼀个特殊的函数来进⾏模型的保存
+            # For zero stage 3, each gpu only has a part of the model, so we need a special save function
+            #save_zero_three_model(model,
+            #                      args.global_rank,
+            #                      args.output_dir,
+            #                      zero_stage=args.zero_stage)
+
+            save_zero_three_model_safetensors(model,
+                                              torch.distributed.get_rank(),
+                                              "./output/final_model",
+                                              zero_stage=args.zero_stage,
+                                              lora_alpha=args.lora_dim,
+                                              merge_lora=True)
+            
+            save_model_config_and_tokenizer(
+                model.module if hasattr(model, 'module') else model,
+                torch.distributed.get_rank(),
+                "./output/final_model",
+                base_model_path="workspace/Qwen3-4B"
+            )
+
+if __name__ == "__main__":
+    main()
+    
+
+'''
+import argparse
+import math
+import sys
+sys.path.append("/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat")
+import torch
+from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
+from torch.utils.data.distributed import DistributedSampler
+
+from transformers import (
+    AutoModelForCausalLM,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+
+import deepspeed
+from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
+from deepspeed import get_accelerator
+
+from dschat.utils.data.data_utils import create_prompt_dataset
+from dschat.utils.utils import print_rank_0, to_device, save_hf_format, set_random_seed, get_all_reduce_mean, get_optimizer_grouped_parameters, save_zero_three_model, load_hf_tokenizer, save_hf_format_safetensors
+from dschat.utils.ds_utils import get_train_ds_config
+from dschat.utils.module.lora import convert_linear_layer_to_lora, convert_lora_to_linear_layer, only_optimize_lora_parameters, make_model_gradient_checkpointing_compatible
+from dschat.utils.model.model_utils import create_hf_model, causal_lm_model_to_fp32_loss
+from dschat.utils.perf import print_throughput
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description=
+        "Finetune a transformers model on a causal language modeling task")
+    # 默认的数据集我们不⽤, data_path传参进来的是垂直领域的个性化数据集
+    parser.add_argument('--data_path',
+                        nargs='*',
+                        default=['Dahoas/rm-static'],
+                        help='Path to the training dataset. Accepted format:'
+                        '1) a single data path, 2) multiple datasets in the'
+                        'form: dataset1-path dataset2-path ...')
+    parser.add_argument('--data_split',
+                        type=str,
+                        default='6,2,2',
+                        help='Comma-separated list of proportions for training'
+                        'phase 1, 2, and 3 data. For example the split `6,2,2`'
+                        'will use 60%% of data for phase 1, 20%% for phase 2'
+                        'and 20%% for phase 3.')
+    parser.add_argument(
+        '--sft_only_data_path',
+        nargs='*',
+        default=[],
+        help='Path to the dataset for only using in SFT phase.')
+    parser.add_argument(
+        '--data_output_path',
+        type=str,
+        default='/tmp/data_files/',
+        help=
+        'Where to store the data-related files such as shuffle index. This needs to be on a local storage of a node (not on a shared storage)'
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help=
+        "Path to pretrained model or model identifier from huggingface.co/models.",
+        required=True,
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--max_seq_len",
+        type=int,
+        default=512,
+        help="The maximum sequence length.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-3,
+        help=
+        "Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay",
+                        type=float,
+                        default=0.,
+                        help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs",
+                        type=int,
+                        default=1,
+                        help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help=
+        "Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="cosine",
+        help="The scheduler type to use.",
+        choices=[
+            "linear", "cosine", "cosine_with_restarts", "polynomial",
+            "constant", "constant_with_warmup"
+        ],
+    )
+    parser.add_argument(
+        "--num_warmup_steps",
+        type=int,
+        default=0,
+        help="Number of steps for the warmup in the lr scheduler.")
+    parser.add_argument("--output_dir",
+                        type=str,
+                        default=None,
+                        help="Where to store the model.")
+    parser.add_argument("--seed",
+                        type=int,
+                        default=1234,
+                        help="A seed for reproducible training.")
+    parser.add_argument("--local_rank",
+                        type=int,
+                        default=-1,
+                        help="local_rank for distributed training on gpus")
+    parser.add_argument('--gradient_checkpointing',
+                        action='store_true',
+                        help='Enable HF gradient checkpointing for model.')
+    parser.add_argument(
+        "--dropout",
+        type=float,
+        default=None,
+        help="If dropout configured, use it. "
+        "Otherwise, keep the default dropout configuration of the model.")
+    # deepspeed features
+    parser.add_argument('--offload',
+                        action='store_true',
+                        help='Enable ZeRO Offload techniques.')
+    parser.add_argument('--dtype',
+                        type=str,
+                        default='fp16',
+                        choices=['fp16', 'bf16'],
+                        help='Training data type')
+    parser.add_argument(
+        '--zero_stage',
+        type=int,
+        default=0,
+        help='ZeRO optimization stage for Actor model (and clones).')
+    ## LoRA for efficient training setting
+    parser.add_argument("--lora_dim",
+                        type=int,
+                        default=0,
+                        help="If > 0, use LoRA for efficient training.")
+    parser.add_argument("--lora_module_name",
+                        type=str,
+                        default="decoder.layers.",
+                        help="The scope of LoRA.")
+    parser.add_argument('--only_optimize_lora',
+                        action='store_true',
+                        help='Only optimize the LoRA parameters.')
+    parser.add_argument(
+        "--lora_learning_rate",
+        type=float,
+        default=5e-4,
+        help=
+        "Initial LoRA learning rate (after the potential warmup period) to use."
+    )
+    ## low precision
+    parser.add_argument(
+        '--compute_fp32_loss',
+        action='store_true',
+        help='Relevant for low precision dtypes (fp16, bf16, etc.). '
+        'If specified, loss is calculated in fp32.')
+    ## Tensorboard logging
+    parser.add_argument('--enable_tensorboard',
+                        action='store_true',
+                        help='Enable tensorboard logging')
+    parser.add_argument('--tensorboard_path',
+                        type=str,
+                        default="step1_tensorboard")
+    ## Tokenizer
+    parser.add_argument(
+        "--add_eot_token",
+        action='store_true',
+        help="Add `eot_token` as additional special token to tokenizer")
+    parser.add_argument(
+        "--eot_token",
+        type=str,
+        default="<|endoftext|>",
+        help="Specify the format of the `eot_token`",
+    )
+    ## Print loss
+    parser.add_argument('--print_loss',
+                        action='store_true',
+                        help='Prints loss at each step.')
+    # 此处是所有超参数和训练参数的设置位置
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args()
+
+    return args
+
+
+def main():
+    args = parse_args()
+
+    if args.local_rank == -1:
+        device = torch.device(get_accelerator().device_name())
+    else:
+        get_accelerator().set_device(args.local_rank)
+        device = torch.device(get_accelerator().device_name(), args.local_rank)
+        # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
+        # torch.distributed.init_process_group(backend='nccl')
+        deepspeed.init_distributed()
+
+    args.global_rank = torch.distributed.get_rank()
+
+    ds_config = get_train_ds_config(offload=args.offload,
+                                    dtype=args.dtype,
+                                    stage=args.zero_stage,
+                                    enable_tensorboard=args.enable_tensorboard,
+                                    tb_path=args.tensorboard_path,
+                                    tb_name="step1_model")
+    ds_config[
+        'train_micro_batch_size_per_gpu'] = args.per_device_train_batch_size
+    ds_config[
+        'train_batch_size'] = args.per_device_train_batch_size * torch.distributed.get_world_size(
+        ) * args.gradient_accumulation_steps
+
+    # If passed along, set the training seed now.
+    set_random_seed(args.seed)
+
+    torch.distributed.barrier()
+
+    # 实例化tokenizer和model
+    # load_hf_tokenizer will get the correct tokenizer and set padding tokens based on the model family
+    additional_special_tokens = args.eot_token if args.add_eot_token else None
+    tokenizer = load_hf_tokenizer(args.model_name_or_path,
+                                  fast_tokenizer=True,
+                                  add_special_tokens=additional_special_tokens)
+
+    model = create_hf_model(AutoModelForCausalLM,
+                            args.model_name_or_path,
+                            tokenizer,
+                            ds_config,
+                            dropout=args.dropout)
+
+    if args.compute_fp32_loss:
+        print_rank_0(
+            f"Using model {model.__class__.__name__} with loss in fp32",
+            args.global_rank)
+        causal_lm_model_to_fp32_loss(model)
+
+    # 设置LoRA微调
+    if args.lora_dim > 0:
+        model = convert_linear_layer_to_lora(model, args.lora_module_name,
+                                             args.lora_dim)
+        if args.only_optimize_lora:
+            model = only_optimize_lora_parameters(model)
+            model = make_model_gradient_checkpointing_compatible(model)
+
+    # 准备训练数据, 注意当前处于第⼀阶段 SFT
+    # Prepare the data
+    train_phase = 1
+    print('args: ', args)
+    print('data_path: ', args.data_path)
+    train_dataset, eval_dataset = create_prompt_dataset(
+        args.local_rank,
+        args.data_path,
+        args.data_split,
+        args.data_output_path,
+        train_phase,
+        args.seed,
+        tokenizer,
+        args.max_seq_len,
+        end_of_conversation_token=tokenizer.eos_token,
+        sft_only_data_path=args.sft_only_data_path)
+    # DataLoaders creation:
+    if args.local_rank == -1:
+        train_sampler = RandomSampler(train_dataset)
+        eval_sampler = SequentialSampler(eval_dataset)
+    else:
+        train_sampler = DistributedSampler(train_dataset)
+        eval_sampler = DistributedSampler(eval_dataset)
+    train_dataloader = DataLoader(train_dataset,
+                                  collate_fn=default_data_collator,
+                                  sampler=train_sampler,
+                                  batch_size=args.per_device_train_batch_size)
+    eval_dataloader = DataLoader(eval_dataset,
+                                 collate_fn=default_data_collator,
+                                 sampler=eval_sampler,
+                                 batch_size=args.per_device_eval_batch_size)
+
+    # main内部定义的评估函数
+    def evaluation(model, eval_dataloader):
+        model.eval()
+        losses = 0
+        for step, batch in enumerate(eval_dataloader):
+            batch = to_device(batch, device)
+            with torch.no_grad():
+                outputs = model(**batch)
+
+            loss = outputs.loss
+            losses += loss.float()
+        losses = losses / (step + 1)
+        try:
+            losses = get_all_reduce_mean(losses)
+        except:
+            pass
+        try:
+            perplexity = torch.exp(losses).item()
+        except OverflowError:
+            perplexity = float("inf")
+        return perplexity, losses.item()
+
+    # 采⽤分组优化参数的优化器策略
+    # Split weights in two groups, one with weight decay and the other not.
+    optimizer_grouped_parameters = get_optimizer_grouped_parameters(
+        model, args.weight_decay, args.lora_learning_rate)
+
+    AdamOptimizer = DeepSpeedCPUAdam if args.offload else FusedAdam
+    optimizer = AdamOptimizer(optimizer_grouped_parameters,
+                              lr=args.learning_rate,
+                              betas=(0.9, 0.95))
+
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader) / args.gradient_accumulation_steps)
+    lr_scheduler = get_scheduler(
+        name=args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=args.num_warmup_steps,
+        num_training_steps=args.num_train_epochs * num_update_steps_per_epoch,
+    )
+
+    # 采⽤deepspeed对相关组件进⾏封装, 本质上是进⾏加速优化
+    model, optimizer, _, lr_scheduler = deepspeed.initialize(
+        model=model,
+        optimizer=optimizer,
+        args=args,
+        config=ds_config,
+        lr_scheduler=lr_scheduler,
+        dist_init_required=True)
+
+    if args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+
+    # 开始训练, 打印⼀些关键信息
+    # Train!
+    print_rank_0("***** Running training *****", args.global_rank)
+    print_rank_0(
+        f"***** Evaluating perplexity, Epoch {0}/{args.num_train_epochs} *****",
+        args.global_rank)
+    perplexity, eval_loss = evaluation(model, eval_dataloader)
+    print_rank_0(f"ppl: {perplexity}, loss: {eval_loss}", args.global_rank)
+
+    # 经典的双重for循环训练模式
+    for epoch in range(args.num_train_epochs):
+        print_rank_0(
+            f"Beginning of Epoch {epoch+1}/{args.num_train_epochs}, Total Micro Batches {len(train_dataloader)}",
+            args.global_rank)
+        # 将模型设置为训练模式
+        model.train()
+        import time
+        for step, batch in enumerate(train_dataloader):
+            start = time.time()
+            batch = to_device(batch, device)
+            # 模型的前向传播计算, 并取到损失值loss
+            outputs = model(**batch, use_cache=False)
+            loss = outputs.loss
+            if args.print_loss:
+                print(
+                    f"Epoch: {epoch}, Step: {step}, Rank: {torch.distributed.get_rank()}, loss = {loss}"
+                )
+            # 反向传播, "⽼三样"
+            model.backward(loss)
+            model.step()
+            end = time.time()
+            if torch.distributed.get_rank() == 0:
+                print_throughput(model.model, args, end - start,
+                                 args.global_rank)
+
+        # 在验证集上进⾏评估, 获取困惑度
+        # Evaluate perplexity on the validation set.
+        print_rank_0(
+            f"***** Evaluating perplexity, Epoch {epoch+1}/{args.num_train_epochs} *****",
+            args.global_rank)
+        perplexity, eval_loss = evaluation(model, eval_dataloader)
+        print_rank_0(f"ppl: {perplexity}, loss: {eval_loss}", args.global_rank)
+        model.tput_timer.update_epoch_count()
+
+    if args.output_dir is not None:
+        print_rank_0('saving the final model ...', args.global_rank)
+        model = convert_lora_to_linear_layer(model)
+
+        if args.global_rank == 0:
+            # save_hf_format(model, tokenizer, args)
+            # 因为我们项⽬中需要训练的是Qwen3⼤模型, 需要保存成safetensor的格式
+            save_hf_format_safetensors(model, tokenizer, args)
+
+        if args.zero_stage == 3:
+            # 在zero_stage==3时, 每⼀个GPU只包含model的⼀部分, 因此需要⼀个特殊的函数来进⾏模型的保存
+            # For zero stage 3, each gpu only has a part of the model, so we need a special save function
+            #save_zero_three_model(model,
+            #                      args.global_rank,
+            #                      args.output_dir,
+            #                      zero_stage=args.zero_stage)
+
+            save_zero_three_model_safetensors(model,
+                                              torch.distributed.get_rank(),
+                                              "./output/final_model",
+                                              zero_stage=args.zero_stage,
+                                              lora_alpha=args.lora_dim,
+                                              merge_lora=True)
+            
+            save_model_config_and_tokenizer(
+                model.module if hasattr(model, 'module') else model,
+                torch.distributed.get_rank(),
+                "./output/final_model",
+                base_model_path="workspace/Qwen3-4B"
+            )
+
+if __name__ == "__main__":
+    main()
+'''

SFT-EN-01-29-2026/code/model_utils.py

@@ -0,0 +1,168 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import math
+import time
+import torch
+from transformers import (
+    AutoConfig,
+    AutoModel,
+)
+from huggingface_hub import snapshot_download
+from transformers.integrations import HfDeepSpeedConfig
+
+from .reward_model import RewardModel
+from ..utils import load_state_dict_into_model
+
+
+def configure_dropout(model_config, dropout):
+    if dropout is not None:
+        for key in ('dropout', 'attention_dropout', 'hidden_dropout',
+                    'activation_dropout'):
+            if hasattr(model_config, key):
+                print(f"Setting model_config.{key} to {dropout}")
+                setattr(model_config, key, dropout)
+
+
+def causal_lm_model_to_fp32_loss(model):
+    """ Convert CausalLM model to calculate loss in fp32 """
+
+    def causal_lm_forward(
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **deprecated_arguments,
+    ):
+        kwargs = dict() if model.config.model_type == "llama" else dict(
+            head_mask=head_mask)
+        output = model.__original_forward__(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            labels=None,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs)
+
+        return_dict = isinstance(output, dict)
+        lm_logits = output.logits if return_dict else output[0]
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(lm_logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].float().contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            batch_size, seq_length, vocab_size = shift_logits.shape
+            # Flatten the tokens
+            loss_fct = torch.nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(batch_size * seq_length, vocab_size),
+                shift_labels.view(batch_size * seq_length))
+
+        if not return_dict:
+            # re-pack output with fp32 loss
+            return ((loss, ) + output) if loss is not None else output
+
+        output.loss = loss
+        return output
+
+    model.__original_forward__ = model.forward
+    model.forward = causal_lm_forward
+
+
+def create_hf_model(model_class,
+                    model_name_or_path,
+                    tokenizer,
+                    ds_config=None,
+                    rlhf_training=False,
+                    dropout=None):
+    model_config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
+    configure_dropout(model_config, dropout)
+
+    # Note: dschf is defined in function scope to avoid global effects
+    # https://huggingface.co/docs/transformers/main_classes/deepspeed#nontrainer-deepspeed-integration
+    if ds_config is not None and ds_config["zero_optimization"]["stage"] == 3:
+        dschf = HfDeepSpeedConfig(ds_config)
+    else:
+        dschf = None
+    if rlhf_training:
+        # the weight loading is handled by create critic model
+        with no_init_weights():
+            model = model_class.from_config(model_config)
+    else:
+        from transformers import AutoModelForCausalLM as _AutoModel
+        model = _AutoModel.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=True,
+            torch_dtype="auto",
+            device_map=None)
+
+    model.config.end_token_id = tokenizer.eos_token_id
+    model.config.pad_token_id = model.config.eos_token_id
+    model.resize_token_embeddings(int(
+        8 *
+        math.ceil(len(tokenizer) / 8.0)))  # make the vocab size multiple of 8
+
+    return model
+
+def create_critic_model(model_name_or_path,
+                       tokenizer,
+                       ds_config,
+                       num_padding_at_beginning=0,
+                       rlhf_training=False,
+                       disable_dropout=False,
+                       zero_stage=0):
+    start = time.time()
+    # 创建critic_model, 本质上也是调用上面的create_hf_model()函数
+    critic_model = create_hf_model(AutoModel, model_name_or_path, tokenizer,
+                                  ds_config, rlhf_training, disable_dropout)
+    end = time.time()
+    # 单独运行第二阶段训练Reward Model的评估代码run_eval.sh时, 可能有报错, 可以暂时先注释下面两行即可
+    if torch.distributed.get_rank() == 0:
+        print(f"> Creating model from_config took {end - start} seconds")
+    
+    critic_model = RewardModel(critic_model,
+                              tokenizer,
+                              num_padding_at_beginning=num_padding_at_beginning)
+    
+    if rlhf_training:
+        # load critic model from checkpoint
+        if not os.path.isdir(model_name_or_path):
+            model_name_or_path = snapshot_download(model_name_or_path)
+        model_ckpt_path = os.path.join(model_name_or_path, 'pytorch_model.bin')
+        assert os.path.exists(model_ckpt_path), f"Cannot find model checkpoint at {model_ckpt_path}"
+        
+        start = time.time()
+        model_ckpt_state_dict = torch.load(model_ckpt_path, map_location='cpu')
+        end = time.time()
+        # 单独运行第二阶段训练Reward Model的评估代码run_eval.sh时, 有报错, 可以暂时先注释下面两行即可
+        if torch.distributed.get_rank() == 0:
+            print(f"> torch.load took {end - start} seconds")
+        
+        # load critic model from checkpoint with zero-stage 3 compatibility
+        # this functionality may be moved to DS checkpoint load API in future
+        start = time.time()
+        load_state_dict_into_model(critic_model,
+                                  model_ckpt_state_dict,
+                                  "",
+                                  zero_stage=zero_stage)
+        end = time.time()
+        # 单独运行第二阶段训练Reward Model的评估代码run_eval.sh时, 有报错, 可以暂时先注释下面两行即可
+        if torch.distributed.get_rank() == 0:
+            print(f"> Loading model state dict took {end - start} seconds")
+    
+    return critic_model

SFT-EN-01-29-2026/code/prompt_eval.py

@@ -0,0 +1,146 @@
+import argparse
+import logging
+import torch
+
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description="Eval baseline vs finetuned SFT model (clean compare)")
+    p.add_argument("--model_name_or_path_baseline", type=str, required=True)
+    p.add_argument("--model_name_or_path_finetune", type=str, required=True)
+    p.add_argument("--max_new_tokens", type=int, default=200)
+    p.add_argument("--language", type=str, default="English", choices=["English", "Chinese"])
+    p.add_argument("--device", type=str, default=None, help="cuda / cpu. default: auto")
+    return p.parse_args()
+
+
+def load_tokenizer(path: str):
+    tok = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    return tok
+
+
+class StopOnSubsequence(StoppingCriteria):
+    def __init__(self, stop_token_seqs):
+        super().__init__()
+        self.stop_token_seqs = stop_token_seqs  # List[List[int]]
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        seq = input_ids[0].tolist()
+        for stop_seq in self.stop_token_seqs:
+            if len(stop_seq) == 0:
+                continue
+            if len(seq) >= len(stop_seq) and seq[-len(stop_seq):] == stop_seq:
+                return True
+        return False
+
+
+def build_stopping_criteria(tokenizer):
+    stop_strings = ["\nHuman:", "\nAssistant:", "Human:", "Assistant:"]
+    stop_token_seqs = [tokenizer.encode(s, add_special_tokens=False) for s in stop_strings]
+    return StoppingCriteriaList([StopOnSubsequence(stop_token_seqs)])
+
+
+def post_trim(text: str):
+    markers = ["\nHuman:", "\nAssistant:", "Human:", "Assistant:"]
+    cut = None
+    for m in markers:
+        idx = text.find(m)
+        if idx != -1:
+            cut = idx if cut is None else min(cut, idx)
+    if cut is not None:
+        text = text[:cut]
+    return text.strip()
+
+
+def generate_greedy(model, tokenizer, prompt, device, max_new_tokens=200):
+    enc = tokenizer(prompt, return_tensors="pt", padding=False, truncation=True, return_attention_mask=True)
+    input_ids = enc["input_ids"].to(device)
+    attention_mask = enc["attention_mask"].to(device)
+
+    stopping_criteria = build_stopping_criteria(tokenizer)
+
+    with torch.no_grad():
+        gen = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            temperature=None,
+            top_p=None,
+            top_k=None,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            stopping_criteria=stopping_criteria,
+        )
+
+    new_tokens = gen[0][input_ids.shape[-1]:]
+    out = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    return post_trim(out)
+
+
+def load_model(path: str, device: torch.device):
+    model = AutoModelForCausalLM.from_pretrained(
+        path,
+        trust_remote_code=True,
+        dtype=torch.bfloat16,
+    )
+    return model.to(device).eval()
+
+
+def get_prompts(language: str):
+    if language == "English":
+        return [
+            "Human: My father was just diagnosed with diabetes a few days ago at the hospital. He is 60 years old with a blood sugar level of 10. What can he eat to improve his condition? Assistant:",
+            "Human: What is hemorrhoid prolapse? What should I do about it? What are the dangers? How should it be treated? Assistant:",
+            "Human: My grandmother is around 70 years old and has had high blood pressure for many years. Recently she has nosebleeds almost every day. A few days ago her submandibular lymph nodes were painful, and the hospital said there might be something wrong with her blood. Could this be leukemia? Assistant:",
+            "Human: My wisdom tooth is inflamed. Yesterday the dentist packed it with medicine but now it is swollen again. What should I do? Assistant:",
+            "Human: These past two days my child's nose seems to be blocked. When lying flat, breathing is difficult, but it gets better when picked up. Sometimes there's also coughing. What should I do? Assistant:",
+            "Human: Four days after intercourse, I tested positive for pregnancy. Ultrasound showed a gestational sac of 15.8mm. Is this pregnancy from this recent intercourse, or was I already pregnant before? Assistant:",
+        ]
+    else:
+        return [
+            "Human: 爸爸前几天检查出糖尿病，60岁血糖10，吃什么能好转? Assistant:",
+            "Human: 什么是痔疮脱出?怎么治疗? Assistant:",
+        ]
+
+
+def main():
+    args = parse_args()
+    device = torch.device("cuda" if (args.device is None and torch.cuda.is_available()) else (args.device or "cpu"))
+
+    tok_base = load_tokenizer(args.model_name_or_path_baseline)
+    tok_ft = load_tokenizer(args.model_name_or_path_finetune)
+
+    print("Loading baseline model...")
+    model_base = load_model(args.model_name_or_path_baseline, device)
+
+    print("Loading finetuned model...")
+    model_ft = load_model(args.model_name_or_path_finetune, device)
+
+    prompts = get_prompts(args.language)
+
+    for i, prompt in enumerate(prompts):
+        print("\n" + "=" * 60)
+        print(f"Prompt {i+1}: {prompt[:80]}...")
+
+        print("\n=== Baseline ===")
+        out_base = generate_greedy(model_base, tok_base, prompt, device, args.max_new_tokens)
+        print(out_base if out_base else "(empty)")
+
+        print("\n=== Finetuned ===")
+        out_ft = generate_greedy(model_ft, tok_ft, prompt, device, args.max_new_tokens)
+        print(out_ft if out_ft else "(empty)")
+
+
+if __name__ == "__main__":
+    main()

SFT-EN-01-29-2026/code/raw_datasets.py

@@ -0,0 +1,828 @@
+# Copyright (c) Microsoft Corporation.
+from datasets import DatasetDict
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+# DeepSpeed Team
+from datasets import load_dataset, load_from_disk
+from torch.utils.data import Subset
+import re
+
+
+# The template prompt dataset class that all new dataset porting needs to
+# follow in order to have a unified API and unified data format.
+class PromptRawDataset(object):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        self.output_path = output_path
+        self.seed = seed
+        self.local_rank = local_rank
+        #if os.path.exists(dataset_name):
+        #    self.raw_datasets = load_from_disk(dataset_name)
+        if not dataset_name == 'local/jsonfile':
+            #self.raw_datasets = load_dataset(dataset_name)
+            self.raw_datasets = None
+
+    def get_train_data(self):
+        return
+
+    def get_eval_data(self):
+        return
+
+    # The prompt should be in the format of: " Human: " + actual_prompt_sentence + " Assistant:"
+    def get_prompt(self, sample):
+        return
+
+    # The chosen response should be in the format of: " " + actual_response_sentence
+    def get_chosen(self, sample):
+        return
+
+    # The rejected response should be in the format of: " " + actual_response_sentence
+    # If the dataset does not have rejected response, return None
+    def get_rejected(self, sample):
+        return
+
+    def get_prompt_and_chosen(self, sample):
+        return
+
+    def get_prompt_and_rejected(self, sample):
+        return
+
+
+# English dataset
+class DahoasRmstaticDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Dahoas/rm-static"
+        self.dataset_name_clean = "Dahoas_rm_static"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["test"]
+
+    def get_prompt(self, sample):
+        return sample['prompt']
+
+    def get_chosen(self, sample):
+        return sample['chosen']
+
+    def get_rejected(self, sample):
+        return sample['rejected']
+
+    def get_prompt_and_chosen(self, sample):
+        return sample['prompt'] + sample['chosen']
+
+    def get_prompt_and_rejected(self, sample):
+        return sample['prompt'] + sample['rejected']
+
+
+# English dataset
+class DahoasFullhhrlhfDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Dahoas/full-hh-rlhf"
+        self.dataset_name_clean = "Dahoas_full_hh_rlhf"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["test"]
+
+    def get_prompt(self, sample):
+        return sample['prompt']
+
+    def get_chosen(self, sample):
+        return sample['chosen']
+
+    def get_rejected(self, sample):
+        return sample['rejected']
+
+    def get_prompt_and_chosen(self, sample):
+        return sample['prompt'] + sample['chosen']
+
+    def get_prompt_and_rejected(self, sample):
+        return sample['prompt'] + sample['rejected']
+
+
+# English dataset
+class DahoasSyntheticinstructgptjpairwiseDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Dahoas/synthetic-instruct-gptj-pairwise"
+        self.dataset_name_clean = "Dahoas_synthetic_instruct_gptj_pairwise"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['prompt'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['chosen']
+
+    def get_rejected(self, sample):
+        return " " + sample['rejected']
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['prompt'] + " Assistant: " + sample['chosen']
+
+    def get_prompt_and_rejected(self, sample):
+        return " Human: " + sample['prompt'] + " Assistant: " + sample[
+            'rejected']
+
+
+# English dataset
+class YitingxieRlhfrewarddatasetsDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "yitingxie/rlhf-reward-datasets"
+        self.dataset_name_clean = "yitingxie_rlhf_reward_datasets"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["test"]
+
+    def get_prompt(self, sample):
+        return sample['prompt'] + "Assistant:"
+
+    def get_chosen(self, sample):
+        return sample['chosen'].split("Assistant:")[-1]
+
+    def get_rejected(self, sample):
+        return sample['rejected'].split("Assistant:")[-1]
+
+    def get_prompt_and_chosen(self, sample):
+        return sample['prompt'] + sample['chosen']
+
+    def get_prompt_and_rejected(self, sample):
+        return sample['prompt'] + sample['rejected']
+
+
+# English dataset
+class OpenaiWebgptcomparisonsDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "openai/webgpt_comparisons"
+        self.dataset_name_clean = "openai_webgpt_comparisons"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['question']['full_text'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        if float(sample['score_0']) >= float(sample['score_1']):
+            response = sample['answer_0']
+        else:
+            response = sample['answer_1']
+        # This data has citation square brackets and numbers (e.g., "[1]").
+        # Right now we are not doing browser-assisted finetuning, thus we
+        # remove these citations to avoid confusing the model.
+        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
+        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
+        return " " + response
+
+    def get_rejected(self, sample):
+        if float(sample['score_0']) < float(sample['score_1']):
+            response = sample['answer_0']
+        else:
+            response = sample['answer_1']
+        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
+        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
+        return " " + response
+
+    def get_prompt_and_chosen(self, sample):
+        if float(sample['score_0']) >= float(sample['score_1']):
+            response = sample['answer_0']
+        else:
+            response = sample['answer_1']
+        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
+        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
+        return " Human: " + sample['question'][
+            'full_text'] + " Assistant: " + response
+
+    def get_prompt_and_rejected(self, sample):
+        if float(sample['score_0']) < float(sample['score_1']):
+            response = sample['answer_0']
+        else:
+            response = sample['answer_1']
+        response = re.sub(r" [\(\[].*?[\)\]]", "", response)
+        response = re.sub(r"[\(\[].*?[\)\]]", "", response)
+        return " Human: " + sample['question'][
+            'full_text'] + " Assistant: " + response
+
+
+# English dataset
+class StanfordnlpSHPDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "stanfordnlp/SHP"
+        self.dataset_name_clean = "stanfordnlp_SHP"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["validation"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['history'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        if int(sample["labels"]) == 1:
+            response = sample["human_ref_A"]
+        else:
+            response = sample["human_ref_B"]
+        return " " + response
+
+    def get_rejected(self, sample):
+        if int(sample["labels"]) == 1:
+            response = sample["human_ref_B"]
+        else:
+            response = sample["human_ref_A"]
+        return " " + response
+
+    def get_prompt_and_chosen(self, sample):
+        if int(sample["labels"]) == 1:
+            response = sample["human_ref_A"]
+        else:
+            response = sample["human_ref_B"]
+        return " Human: " + sample['history'] + " Assistant: " + response
+
+    def get_prompt_and_rejected(self, sample):
+        if int(sample["labels"]) == 1:
+            response = sample["human_ref_B"]
+        else:
+            response = sample["human_ref_A"]
+        return " Human: " + sample['history'] + " Assistant: " + response
+
+
+# English dataset
+class PvduySharegptalpacaoavicunaformatDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "pvduy/sharegpt_alpaca_oa_vicuna_format"
+        self.dataset_name_clean = "pvduy_sharegpt_alpaca_oa_vicuna_format"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["test"]
+
+    def get_prompt(self, sample):
+        if sample['prompt'] is not None and len(sample['prompt']) > 0:
+            return sample['prompt'].replace("USER", "Human").replace(
+                "ASSISTANT", "Assistant")
+        return None
+
+    def get_chosen(self, sample):
+        if sample['label'] is not None and len(sample['label']) > 0:
+            return " " + sample['label']
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['prompt'] is not None and sample['label'] is not None and len(
+                sample['prompt']) > 0 and len(sample['label']) > 0:
+            return sample['prompt'].replace("USER", "Human").replace(
+                "ASSISTANT", "Assistant") + " " + sample['label']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+class LocalJsonFileDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name, chat_path):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "local/jsonfile"
+        self.dataset_name_clean = "jsonfile"
+        self.raw_datasets = load_dataset('json',
+                                         data_files={
+                                             "train":
+                                             chat_path + '/data/train.json',
+                                             "eval":
+                                             chat_path + '/data/eval.json'
+                                         })
+
+    def get_train_data(self):
+        if self.raw_datasets['train'] is not None:
+            return self.raw_datasets['train']
+        return None
+
+    def get_eval_data(self):
+        if self.raw_datasets['eval'] is not None:
+            return self.raw_datasets['eval']
+        return None
+
+    # The prompt should be in the format of: " Human: " + actual_prompt_sentence + " Assistant:"
+    def get_prompt(self, sample):
+        if sample['prompt'] is not None:
+            return " " + sample['prompt']
+        return None
+
+    # The chosen response should be in the format of: " " + actual_response_sentence
+    def get_chosen(self, sample):
+        if sample['chosen'] is not None:
+            return " " + sample['chosen']
+        return None
+
+    # The rejected response should be in the format of: " " + actual_response_sentence
+    # If the dataset does not have rejected response, return None
+    def get_rejected(self, sample):
+        if sample['rejected'] is not None:
+            return " " + sample['rejected']
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['prompt'] is not None and sample['chosen'] is not None:
+            return " " + sample['prompt'] + " " + sample['chosen']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        if sample['prompt'] is not None and sample['rejected'] is not None:
+            return " " + sample['prompt'] + " " + sample['rejected']
+        return None
+
+
+# Chinese dataset
+class Wangrui6ZhihuKOLDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "wangrui6/Zhihu-KOL"
+        self.dataset_name_clean = "wangrui6_Zhihu_KOL"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        if sample['INSTRUCTION'] is not None:
+            return " Human: " + sample['INSTRUCTION'] + " Assistant:"
+        return None
+
+    def get_chosen(self, sample):
+        if sample['RESPONSE'] is not None:
+            return " " + sample['RESPONSE']
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['INSTRUCTION'] is not None and sample['RESPONSE'] is not None:
+            return " Human: " + sample[
+                'INSTRUCTION'] + " Assistant: " + sample['RESPONSE']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Chinese dataset
+class CohereMiraclzhqueries2212Dataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Cohere/miracl-zh-queries-22-12"
+        self.dataset_name_clean = "Cohere_miracl_zh_queries_22_12"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["dev"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['query'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['positive_passages'][0]['text']
+
+    def get_rejected(self, sample):
+        return " " + sample['negative_passages'][0]['text']
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['query'] + " Assistant: " + sample[
+            'positive_passages'][0]['text']
+
+    def get_prompt_and_rejected(self, sample):
+        return " Human: " + sample['query'] + " Assistant: " + sample[
+            'negative_passages'][0]['text']
+
+
+# Chinese dataset
+class HelloSimpleAIHC3ChineseDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Hello-SimpleAI/HC3-Chinese"
+        self.dataset_name_clean = "Hello_SimpleAI_HC3_Chinese"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        if sample['question'] is not None:
+            return " Human: " + sample['question'] + " Assistant:"
+        return None
+
+    def get_chosen(self, sample):
+        if sample['human_answers'][0] is not None:
+            return " " + sample['human_answers'][0]
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['question'] is not None and sample['human_answers'][
+                0] is not None:
+            return " Human: " + sample['question'] + " Assistant: " + sample[
+                'human_answers'][0]
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Chinese dataset
+class MkqaChineseDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "mkqa-Chinese"
+        self.dataset_name_clean = "mkqa"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        if sample['queries']['zh_cn'] is not None:
+            return " Human: " + sample['queries']['zh_cn'] + " Assistant:"
+        return None
+
+    def get_chosen(self, sample):
+        if sample['answers']['zh_cn'][0]['text'] is not None:
+            return " " + sample['answers']['zh_cn'][0]['text']
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['queries']['zh_cn'] is not None and sample['answers'][
+                'zh_cn'][0]['text'] is not None:
+            return " Human: " + sample['queries'][
+                'zh_cn'] + " Assistant: " + sample['answers']['zh_cn'][0][
+                    'text']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Japanese dataset
+class MkqaJapaneseDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "mkqa-Japanese"
+        self.dataset_name_clean = "mkqa"
+
+    def get_train_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 0,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_eval_data(self):
+        from .data_utils import get_raw_dataset_split_index
+        dataset = self.raw_datasets["train"]
+        index = get_raw_dataset_split_index(self.local_rank, self.output_path,
+                                            self.dataset_name_clean,
+                                            self.seed, "train_eval", "9,1", 1,
+                                            len(dataset))
+        dataset = Subset(dataset, index)
+        return dataset
+
+    def get_prompt(self, sample):
+        if sample['queries']['ja'] is not None:
+            return " Human: " + sample['queries']['ja'] + " Assistant:"
+        return None
+
+    def get_chosen(self, sample):
+        if sample['answers']['ja'][0]['text'] is not None:
+            return " " + sample['answers']['ja'][0]['text']
+        return None
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        if sample['queries']['ja'] is not None and sample['answers']['ja'][0][
+                'text'] is not None:
+            return " Human: " + sample['queries'][
+                'ja'] + " Assistant: " + sample['answers']['ja'][0]['text']
+        return None
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Japanese dataset
+class CohereMiracljaqueries2212Dataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "Cohere/miracl-ja-queries-22-12"
+        self.dataset_name_clean = "Cohere_miracl_ja_queries_22_12"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["dev"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['query'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['positive_passages'][0]['text']
+
+    def get_rejected(self, sample):
+        return " " + sample['negative_passages'][0]['text']
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['query'] + " Assistant: " + sample[
+            'positive_passages'][0]['text']
+
+    def get_prompt_and_rejected(self, sample):
+        if len(sample['negative_passages']) > 0:
+            return " Human: " + sample['query'] + " Assistant: " + sample[
+                'negative_passages'][0]['text']
+        return None
+
+
+# Japanese dataset
+class LmqgQgjaquadDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "lmqg/qg_jaquad"
+        self.dataset_name_clean = "lmqg_qg_jaquad"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["validation"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['question'] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['sentence']
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['question'] + " Assistant: " + sample[
+            'sentence']
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+
+# Japanese dataset
+class LmqgQagjaquadDataset(PromptRawDataset):
+
+    def __init__(self, output_path, seed, local_rank, dataset_name):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        self.dataset_name = "lmqg/qag_jaquad"
+        self.dataset_name_clean = "lmqg_qag_jaquad"
+
+    def get_train_data(self):
+        return self.raw_datasets["train"]
+
+    def get_eval_data(self):
+        return self.raw_datasets["validation"]
+
+    def get_prompt(self, sample):
+        return " Human: " + sample['questions'][0] + " Assistant:"
+
+    def get_chosen(self, sample):
+        return " " + sample['paragraph']
+
+    def get_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+
+    def get_prompt_and_chosen(self, sample):
+        return " Human: " + sample['questions'][0] + " Assistant: " + sample[
+            'paragraph']
+
+    def get_prompt_and_rejected(self, sample):
+        print(
+            f"Warning: dataset {self.dataset_name} does not include rejected response."
+        )
+        return None
+# CustomDataset: 自定义数据集类，用于训练个性化垂直领域大模型，继承基类PromptRawDataset
+class CustomDataset(PromptRawDataset):
+    def __init__(self, output_path, seed, local_rank, dataset_name, chat_path):
+        super().__init__(output_path, seed, local_rank, dataset_name)
+        # 个性化数据集的名字可以自定义
+        self.dataset_name = "custom"
+        self.dataset_name_clean = "custom"
+        # 设定要读取的数据集所在的绝对路径
+        train_path = chat_path + '/data/train.jsonl'
+        eval_path = chat_path + '/data/dev.jsonl'
+        # 通过DatasetDict的类封装数据, 和load_dataset()函数保持一致.
+        self.raw_datasets = DatasetDict.from_json({'train': train_path, 'eval': eval_path})
+    
+    # 返回训练集数据
+    def get_train_data(self):
+        if self.raw_datasets['train'] is not None:
+            return self.raw_datasets['train']
+        return None
+    
+    # 返回验证集数据
+    def get_eval_data(self):
+        if self.raw_datasets['eval'] is not None:
+            return self.raw_datasets['eval']
+        return None
+    
+    # 构造prompt输入模型的格式: Human: prompt Assistant:
+    def get_prompt(self, sample):
+        if sample['prompt'] is not None:
+            return " Human: " + sample['prompt'] + " Assistant:"
+        return None
+    
+    # 构造chosen输入模型的格式: chosen
+    def get_chosen(self, sample):
+        if sample['chosen'] is not None:
+            return " " + sample['chosen']
+        return None
+    
+    # 构造reject输入模型的格式: reject
+    def get_rejected(self, sample):
+        if sample['rejected'] is not None:
+            return " " + sample['reject']
+        return None
+    
+    # 构造第二阶段训练Reward Model的输入模型格式: Human: prompt Assistant: chosen
+    def get_prompt_and_chosen(self, sample):
+        if sample['prompt'] is not None and sample['chosen'] is not None:
+            return " Human: " + sample['prompt'] + " Assistant: " + sample['chosen']
+        return None
+    
+    # 构造第二阶段训练Reward Model的输入模型格式: Human: prompt Assistant: reject
+    def get_prompt_and_rejected(self, sample):
+        if sample['prompt'] is not None and sample['reject'] is not None:
+            return " Human: " + sample['prompt'] + " Assistant: " + sample['reject']
+        return None

SFT-EN-01-29-2026/code/utils.py

@@ -0,0 +1,384 @@
+# Copyright (c) Microsoft Corporation.
+from safetensors.torch import save_file
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+import os
+import torch
+import random
+import numpy as np
+from transformers import set_seed, AutoTokenizer
+import json
+import deepspeed
+from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+from deepspeed.accelerator import get_accelerator
+import torch.nn as nn
+
+
+def print_rank_0(msg, rank=None):
+    if rank is not None and rank <= 0:
+        print(msg)
+    elif is_rank_0():
+        print(msg)
+
+
+def is_rank_0():
+    """Check whether it is rank 0."""
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == 0:
+            return True
+        else:
+            return False
+    else:
+        return True
+
+
+def to_device(batch, device):
+    output = {}
+    for k, v in batch.items():
+        try:
+            output[k] = v.to(device)
+        except:
+            output[k] = v
+    return output
+
+
+class MovingAverage:
+
+    def __init__(self):
+        self.count = 0
+        self.total = 0
+        self.mean = 0
+
+    def update(self, num):
+        self.total += num
+        self.count += 1
+        self.mean = self.total / self.count
+
+        return self.mean
+
+
+class ExponentialMovingAverage:
+
+    def __init__(self, alpha=0.9):
+        self.alpha = alpha
+        self.ema = None
+
+    def update(self, num):
+        prev_ema = num if self.ema is None else self.ema
+        self.ema = self.alpha * prev_ema + (1.0 - self.alpha) * num
+        return self.ema
+
+    def get(self):
+        return self.ema if self.ema is not None else 0.
+
+
+def get_tokenizer(model_name_or_path, fast_tokenizer=True):
+    if "llama" in model_name_or_path:
+        from transformers.models.llama import LlamaTokenizer
+        tokenizer = LlamaTokenizer.from_pretrained(
+            model_name_or_path, fast_tokenizer=fast_tokenizer)
+        if tokenizer.pad_token is None:
+            # assert tokenizer.eos_token is not None
+            # tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+            tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+            tokenizer.padding_side = 'right'
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path, fast_tokenizer=fast_tokenizer)
+        tokenizer.pad_token = tokenizer.eos_token
+        # make sure tokenizer is right pad in our logic
+        tokenizer.padding_side = 'right'
+    return tokenizer
+
+
+def load_hf_tokenizer(model_name_or_path,
+                      fast_tokenizer=True,
+                      add_special_tokens=None):
+    if os.path.exists(model_name_or_path):
+        # Locally tokenizer loading has some issue, so we need to force download
+        model_json = os.path.join(model_name_or_path, "config.json")
+        if os.path.exists(model_json):
+            model_json_file = json.load(open(model_json))
+            model_name = model_json_file.get("_name_or_path",
+                                             model_name_or_path)
+            tokenizer = get_tokenizer(model_name,
+                                      fast_tokenizer=fast_tokenizer)
+    else:
+        tokenizer = get_tokenizer(model_name_or_path,
+                                  fast_tokenizer=fast_tokenizer)
+
+    if add_special_tokens is not None:
+        add_special_tokens = [add_special_tokens] if isinstance(add_special_tokens, str) \
+            else add_special_tokens
+        tokenizer.add_special_tokens(
+            {'additional_special_tokens': add_special_tokens})
+
+    return tokenizer
+
+def save_hf_format_safetensors(model, tokenizer, args, sub_folder=""):
+    """
+    将模型和分词器保存为 Hugging Face 格式, 并使用 safetensors 保存模型权重.
+    此版本能正确处理共享内存的张量 (如Qwen3的 lm_head 和 embed_tokens).
+    
+    Args:
+        model: 要保存的模型
+        tokenizer: 分词器
+        args: 包含 output_dir 等参数的对象
+        sub_folder (str, optional): 在输出目录下的子文件夹名
+    """
+    # 1: 提取原始模型 (移除 DeepSpeed 或 DataParallel 的包装)
+    model_to_save = model.module if hasattr(model, 'module') else model
+    
+    # 2: 定义输出路径
+    output_dir = os.path.join(args.output_dir, sub_folder)
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 3: 获取模型状态字典
+    state_dict = model_to_save.state_dict()
+    
+    # 4: 处理共享内存的张量, 创建一个新的字典, 其中共享内存的张量将被克隆
+    new_state_dict = {}
+    # 用于追踪已处理过的内存地址, 避免重复克隆同一内存块
+    seen_data_ptrs = {}
+    
+    for key, tensor in state_dict.items():
+        # 检查张量的底层数据指针
+        data_ptr = tensor.data_ptr()
+        
+        if data_ptr in seen_data_ptrs:
+            # 如果这个内存地址已经出现过, 说明是共享内存张量, 需要克隆一份
+            print(f"检测到共享��存张量 '{key}' 与 '{seen_data_ptrs[data_ptr]}' 共享内存, 正在克隆...")
+            # 使用 .clone() 创建一份独立的副本
+            new_state_dict[key] = tensor.clone()
+        else:
+            # 首次遇到的内存地址, 直接存入新字典并记录
+            new_state_dict[key] = tensor
+            seen_data_ptrs[data_ptr] = key
+    
+    # 5: 移除 LoRA 权重 (如果使用了LoRA微调)
+    if hasattr(model_to_save, 'peft_config') or any("lora" in k for k in new_state_dict.keys()):
+        print("检测到LoRA权重, 正在移除...")
+        keys_to_remove = [key for key in new_state_dict.keys() if "lora" in key]
+        for key in keys_to_remove:
+            del new_state_dict[key]
+            print(f"  已移除: {key}")
+    
+    # 6: 使用 safetensors 保存处理后的权重
+    output_safetensors_file = os.path.join(output_dir, "model.safetensors")
+    # 注意: 这里保存的是 new_state_dict, 而不是原始的 state_dict
+    save_file(new_state_dict, output_safetensors_file, metadata={"format": "pt"})
+    print(f"✅ 模型权重已保存至: {output_safetensors_file}")
+    
+    # 7: 保存模型配置
+    output_config_file = os.path.join(output_dir, "config.json")
+    model_to_save.config.to_json_file(output_config_file)
+    print(f"✅ 模型配置已保存至: {output_config_file}")
+    
+    # 8: 保存分词器 (推荐的标准方式)
+    tokenizer.save_pretrained(output_dir)
+    print(f"✅ 分词器文件已保存至: {output_dir}")
+    
+    # 9: 可选: 验证保存的权重可以正确加载
+    print("正在进行快速加载验证...")
+    try:
+        # 从保存的文件加载权重, 检查完整性
+        from safetensors.torch import load_file
+        loaded_tensors = load_file(output_safetensors_file)
+        print(f"✅ 验证通过! 成功加载了 {len(loaded_tensors)} 个张量.")
+    except Exception as e:
+        print(f"⚠ 加载验证时出现警告(可能不影响后续使用): {e}")
+
+
+def save_hf_format(model, tokenizer, args, sub_folder=""):
+    # used to save huggingface format, so we can use it for hf.from_pretrained
+    model_to_save = model.module if hasattr(model, 'module') else model
+    CONFIG_NAME = "config.json"
+    WEIGHTS_NAME = "pytorch_model.bin"
+    output_dir = os.path.join(args.output_dir, sub_folder)
+    os.makedirs(output_dir, exist_ok=True)
+    output_model_file = os.path.join(output_dir, WEIGHTS_NAME)
+    output_config_file = os.path.join(output_dir, CONFIG_NAME)
+    save_dict = model_to_save.state_dict()
+    for key in list(save_dict.keys()):
+        if "lora" in key:
+            del save_dict[key]
+    torch.save(save_dict, output_model_file)
+    model_to_save.config.to_json_file(output_config_file)
+    tokenizer.save_vocabulary(output_dir)
+
+
+def set_random_seed(seed):
+    if seed is not None:
+        set_seed(seed)
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        get_accelerator().manual_seed_all(seed)
+
+
+def get_all_reduce_mean(tensor):
+    torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
+    tensor = tensor / torch.distributed.get_world_size()
+    return tensor
+
+
+# This function is a modified version of code available in the from_pretrained API of HuggingFace Transformers
+# The code is copied and modified from: https://github.com/huggingface/transformers/blob/5ee9693a1c77c617ebc43ef20194b6d3b674318e/src/transformers/modeling_utils.py#L498
+# This function helps load a HF format checkpoint into a DeepSpeed wrapped model that has been sharded using ZeRO Stage 3
+def load_state_dict_into_model(model_to_load=None,
+                               state_dict=None,
+                               start_prefix="",
+                               zero_stage=0):
+
+    # copy state_dict so _load_from_state_dict can modify it
+    metadata = getattr(state_dict, "_metadata", None)
+    state_dict = state_dict.copy()
+    if metadata is not None:
+        state_dict._metadata = metadata
+
+    error_msgs = []
+
+    # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
+    # so we need to apply the function recursively.
+    def load(module: nn.Module, state_dict, prefix=""):
+        local_metadata = {} if metadata is None else metadata.get(
+            prefix[:-1], {})
+        args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)
+        # Parameters of module and children will start with prefix. We can exit early if there are none in this
+        # state_dict
+        if len([key for key in state_dict if key.startswith(prefix)]) > 0:
+            if zero_stage == 3:
+                # In sharded models, each shard has only part of the full state_dict, so only gather
+                # parameters that are in the current state_dict.
+                named_parameters = dict(
+                    module.named_parameters(prefix=prefix[:-1], recurse=False))
+                params_to_gather = [
+                    named_parameters[k] for k in state_dict.keys()
+                    if k in named_parameters
+                ]
+                if len(params_to_gather) > 0:
+                    # because zero3 puts placeholders in model params, this context
+                    # manager gathers (unpartitions) the params of the current layer, then loads from
+                    # the state dict and then re-partitions them again
+                    with deepspeed.zero.GatheredParameters(params_to_gather,
+                                                           modifier_rank=0):
+                        if torch.distributed.get_rank() == 0:
+                            module._load_from_state_dict(*args)
+            else:
+                module._load_from_state_dict(*args)
+
+        for name, child in module._modules.items():
+            if child is not None:
+                load(child, state_dict, prefix + name + ".")
+
+    load(model_to_load, state_dict, prefix=start_prefix)
+    # Delete `state_dict` so it could be collected by GC earlier. Note that `state_dict` is a copy of the argument, so
+    # it's safe to delete it.
+    del state_dict
+
+    return error_msgs
+
+
+def get_optimizer_grouped_parameters(
+    model,
+    weight_decay,
+    lora_lr=5e-4,
+    no_decay_name_list=[
+        "bias", "layer_norm.weight", "layernorm.weight", "norm.weight",
+        "ln_f.weight"
+    ],
+    lora_name_list=["lora_right_weight", "lora_left_weight"],
+):
+    optimizer_grouped_parameters = [
+        {
+            "params": [
+                p for n, p in model.named_parameters()
+                if (not any(nd in n.lower() for nd in no_decay_name_list)
+                    and p.requires_grad and not any(nd in n.lower()
+                                                    for nd in lora_name_list))
+            ],
+            "weight_decay":
+            weight_decay,
+        },
+        {
+            "params": [
+                p for n, p in model.named_parameters()
+                if (not any(nd in n.lower() for nd in no_decay_name_list)
+                    and p.requires_grad and any(nd in n.lower()
+                                                for nd in lora_name_list))
+            ],
+            "weight_decay":
+            weight_decay,
+            "lr":
+            lora_lr
+        },
+        {
+            "params": [
+                p for n, p in model.named_parameters()
+                if (any(nd in n.lower()
+                        for nd in no_decay_name_list) and p.requires_grad)
+            ],
+            "weight_decay":
+            0.0,
+        },
+    ]
+
+    non_empty_groups = []
+    for group in optimizer_grouped_parameters:
+        if group["params"]:
+            non_empty_groups.append(group)
+    return non_empty_groups
+
+
+def _z3_params_to_fetch(param_list):
+    return [
+        p for p in param_list
+        if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE
+    ]
+
+
+def moving_average(model, model_ema, beta=0.992, device=None, zero_stage=0):
+    zero_stage_3 = (zero_stage == 3)
+    with torch.no_grad():
+        for param, param_ema in zip(model.parameters(),
+                                    model_ema.parameters()):
+            # TODO: use prefiltering for efficiency
+            params_to_fetch = _z3_params_to_fetch([param, param_ema
+                                                   ]) if zero_stage_3 else []
+            should_gather_param = len(params_to_fetch) > 0
+            with deepspeed.zero.GatheredParameters(
+                    params_to_fetch, enabled=should_gather_param):
+                data = param.data
+                if device is not None:
+                    data = data.to(device)
+                param_ema.data.copy_(torch.lerp(data, param_ema.data, beta))
+
+
+def save_zero_three_model(model_ema, global_rank, save_dir, zero_stage=0):
+    zero_stage_3 = (zero_stage == 3)
+    os.makedirs(save_dir, exist_ok=True)
+    WEIGHTS_NAME = "pytorch_model.bin"
+    output_model_file = os.path.join(save_dir, WEIGHTS_NAME)
+
+    model_to_save = model_ema.module if hasattr(model_ema,
+                                                'module') else model_ema
+    if not zero_stage_3:
+        if global_rank == 0:
+            torch.save(model_to_save.state_dict(), output_model_file)
+    else:
+        output_state_dict = {}
+        for k, v in model_to_save.named_parameters():
+
+            if hasattr(v, 'ds_id'):
+                with deepspeed.zero.GatheredParameters(_z3_params_to_fetch([v
+                                                                            ]),
+                                                       enabled=zero_stage_3):
+                    v_p = v.data.cpu()
+            else:
+                v_p = v.cpu()
+            if global_rank == 0 and "lora" not in k:
+                output_state_dict[k] = v_p
+        if global_rank == 0:
+            torch.save(output_state_dict, output_model_file)
+        del output_state_dict

SFT-EN-01-29-2026/data/train.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de56cd90e05715d0521515aa4a90d718d3e0da27d49970ff0a83136652066906
+size 25584972

SFT-EN-01-29-2026/model/chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

SFT-EN-01-29-2026/model/config.json

@@ -0,0 +1,72 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "end_token_id": 151645,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9728,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 36,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "pad_token_id": 151645,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.0.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151672
+}

SFT-EN-01-29-2026/model/ds_tensorboard_logs/step1_model_tensorboard/events.out.tfevents.1769725308.209-20-158-64.30075.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bec6d2eb21f2a317beea31714d84e4c3fe34ba0c62365be1e2dc9ea98806cd55
+size 204

SFT-EN-01-29-2026/model/ds_tensorboard_logs/step1_model_tensorboard/events.out.tfevents.1769725536.209-20-158-64.31271.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:543d7327217546fd37d3fee5322ec0f7ccbb530b17e65a16f14250a52daa3a4a
+size 1448

SFT-EN-01-29-2026/model/ds_tensorboard_logs/step1_model_tensorboard/events.out.tfevents.1769726189.209-20-158-64.32221.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d996a4ea809b4cd08593bf3563f579d16402dfa2a1f23aa03d3110beda00a0e
+size 37198

SFT-EN-01-29-2026/model/ds_tensorboard_logs/step1_model_tensorboard/events.out.tfevents.1769727296.209-20-158-64.32989.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17e264bbafe4a0f20c4f2b0df948cee4cc696f181bd12f2530404e8c71c06444
+size 37198

SFT-EN-01-29-2026/model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9956493572a7ae7ff86699c23789cba8d31a38d0a2d6333177d846b9d9cade23
+size 8820191160

SFT-EN-01-29-2026/model/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650

SFT-EN-01-29-2026/model/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "fast_tokenizer": true,
+  "is_local": true,
+  "model_max_length": 131072,
+  "pad_token": "<|im_end|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

SFT-EN-01-29-2026/model/training.log

@@ -0,0 +1,317 @@
+/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4
+  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"
+[2026-01-29 22:24:38,868] [WARNING] [runner.py:232:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.
+[2026-01-29 22:24:38,868] [INFO] [runner.py:630:main] cmd = /usr/bin/python3 -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbMF19 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None --log_level=info main.py --model_name_or_path /workspace/Qwen3-4B --data_path /home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl --weight_decay 0.1 --dropout 0.0 --gradient_accumulation_steps 8 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --zero_stage 3 --offload --dtype bf16 --enable_tensorboard --tensorboard_path ./output_sft_en --deepspeed --output_dir ./output_sft_en
+/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4
+  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"
+[2026-01-29 22:24:45,395] [INFO] [launch.py:162:main] WORLD INFO DICT: {'localhost': [0]}
+[2026-01-29 22:24:45,396] [INFO] [launch.py:168:main] nnodes=1, num_local_procs=1, node_rank=0
+[2026-01-29 22:24:45,396] [INFO] [launch.py:179:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0]})
+[2026-01-29 22:24:45,396] [INFO] [launch.py:180:main] dist_world_size=1
+[2026-01-29 22:24:45,396] [INFO] [launch.py:184:main] Setting CUDA_VISIBLE_DEVICES=0
+[2026-01-29 22:24:45,398] [INFO] [launch.py:272:main] process 31271 spawned with command: ['/usr/bin/python3', '-u', 'main.py', '--local_rank=0', '--model_name_or_path', '/workspace/Qwen3-4B', '--data_path', '/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl', '--weight_decay', '0.1', '--dropout', '0.0', '--gradient_accumulation_steps', '8', '--per_device_train_batch_size', '1', '--per_device_eval_batch_size', '1', '--zero_stage', '3', '--offload', '--dtype', 'bf16', '--enable_tensorboard', '--tensorboard_path', './output_sft_en', '--deepspeed', '--output_dir', './output_sft_en']
+/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4
+  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"
+[rank0]:[W129 22:24:52.444107661 ProcessGroupNCCL.cpp:4715] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 as device used by this process is currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. You can pecify device_id in init_process_group() to force use of a particular device.
+Setting model_config.attention_dropout to 0.0
+args:  Namespace(data_path=['/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl'], data_split='6,2,2', sft_only_data_path=[], data_output_path='/tmp/data_files/', model_name_or_path='/workspace/Qwen3-4B', per_device_train_batch_size=1, per_device_eval_batch_size=1, max_seq_len=512, learning_rate=0.001, weight_decay=0.1, num_train_epochs=1, gradient_accumulation_steps=8, lr_scheduler_type=<SchedulerType.COSINE: 'cosine'>, num_warmup_steps=0, output_dir='./output_sft_en', seed=1234, local_rank=0, gradient_checkpointing=False, dropout=0.0, offload=True, dtype='bf16', zero_stage=3, lora_dim=0, lora_module_name='decoder.layers.', only_optimize_lora=False, lora_learning_rate=0.0005, compute_fp32_loss=False, enable_tensorboard=True, tensorboard_path='./output_sft_en', add_eot_token=False, eot_token='<|endoftext|>', print_loss=False, deepspeed=True, deepspeed_config=None, deepscale=False, deepscale_config=None, global_rank=0)
+data_path:  ['/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl']
+/usr/lib/python3/dist-packages/torch/utils/cpp_extension.py:2376: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. 
+If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].
+  warnings.warn(
+2026-01-29 22:25:34.798274: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-01-29 22:25:34.808869: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1769725534.821805   31271 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1769725534.825823   31271 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1769725534.835606   31271 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1769725534.835626   31271 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1769725534.835656   31271 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1769725534.835658   31271 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+2026-01-29 22:25:34.838493: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+Stage 3 initialize beginning
+MA 0.72 GB         Max_MA 2.9 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 16.26 GB, percent = 7.4%
+DeepSpeedZeRoOffload initialize [begin]
+MA 0.72 GB         Max_MA 0.72 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 16.25 GB, percent = 7.3%
+Parameter Offload - Persistent parameters statistics: param_count = 145, numel = 196096
+DeepSpeedZeRoOffload initialize [end]
+MA 0.0 GB         Max_MA 0.72 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 16.7 GB, percent = 7.6%
+Before creating fp16 partitions
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 16.7 GB, percent = 7.6%
+After creating fp16 partitions: 5
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 19.89 GB, percent = 9.0%
+Before creating fp32 partitions
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 19.89 GB, percent = 9.0%
+After creating fp32 partitions
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 34.0 GB, percent = 15.4%
+Before initializing optimizer states
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 34.0 GB, percent = 15.4%
+After initializing optimizer states
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 49.09 GB, percent = 22.2%
+After initializing ZeRO optimizer
+MA 0.93 GB         Max_MA 2.38 GB         CA 3.83 GB         Max_CA 4 GB 
+CPU Virtual Memory:  used = 56.32 GB, percent = 25.5%
+***** Running training *****
+Beginning of Epoch 1/1, Total Micro Batches 5400
+Model Parameters: 4.022 B, Latency: 2.91s, TFLOPs: 3.40, Samples/sec: 0.34, Time/seq 2.91s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 3.07s, TFLOPs: 3.22, Samples/sec: 0.33, Time/seq 3.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.34s, TFLOPs: 4.22, Samples/sec: 0.43, Time/seq 2.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.35s, TFLOPs: 4.20, Samples/sec: 0.43, Time/seq 2.35s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.34s, TFLOPs: 4.23, Samples/sec: 0.43, Time/seq 2.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.34s, TFLOPs: 4.22, Samples/sec: 0.43, Time/seq 2.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.33s, TFLOPs: 4.23, Samples/sec: 0.43, Time/seq 2.33s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 6.18s, TFLOPs: 1.60, Samples/sec: 0.16, Time/seq 6.18s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.11s, TFLOPs: 4.69, Samples/sec: 0.47, Time/seq 2.11s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 4.99, Samples/sec: 0.50, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.95, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.96s, TFLOPs: 5.04, Samples/sec: 0.51, Time/seq 1.96s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 4.99, Samples/sec: 0.50, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.19s, TFLOPs: 2.36, Samples/sec: 0.24, Time/seq 4.19s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 5.00, Samples/sec: 0.51, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.97, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.97, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.02, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.92, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.21s, TFLOPs: 2.35, Samples/sec: 0.24, Time/seq 4.21s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.02, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.93, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.02, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 5.00, Samples/sec: 0.51, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.97, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.01, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.24s, TFLOPs: 2.33, Samples/sec: 0.24, Time/seq 4.24s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.39s, TFLOPs: 4.14, Samples/sec: 0.42, Time/seq 2.39s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.36s, TFLOPs: 4.19, Samples/sec: 0.42, Time/seq 2.36s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.31s, TFLOPs: 4.27, Samples/sec: 0.43, Time/seq 2.31s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.96, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.93, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.27s, TFLOPs: 2.31, Samples/sec: 0.23, Time/seq 4.27s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.95s, TFLOPs: 5.06, Samples/sec: 0.51, Time/seq 1.95s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.94s, TFLOPs: 5.09, Samples/sec: 0.52, Time/seq 1.94s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.93s, TFLOPs: 5.12, Samples/sec: 0.52, Time/seq 1.93s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.97, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.27s, TFLOPs: 2.31, Samples/sec: 0.23, Time/seq 4.27s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.13s, TFLOPs: 4.64, Samples/sec: 0.47, Time/seq 2.13s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 4.98, Samples/sec: 0.50, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.98, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.74, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.22s, TFLOPs: 2.34, Samples/sec: 0.24, Time/seq 4.22s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.74, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.75, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.77, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.10s, TFLOPs: 4.71, Samples/sec: 0.48, Time/seq 2.10s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.75, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.30s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.30s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.77, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.25s, TFLOPs: 4.40, Samples/sec: 0.45, Time/seq 2.25s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.29s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.29s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.39s, TFLOPs: 4.13, Samples/sec: 0.42, Time/seq 2.39s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.37s, TFLOPs: 4.17, Samples/sec: 0.42, Time/seq 2.37s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.37s, TFLOPs: 4.18, Samples/sec: 0.42, Time/seq 2.37s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.37s, TFLOPs: 4.17, Samples/sec: 0.42, Time/seq 2.37s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.28s, TFLOPs: 4.33, Samples/sec: 0.44, Time/seq 2.28s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.87, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.28s, TFLOPs: 2.31, Samples/sec: 0.23, Time/seq 4.28s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.10s, TFLOPs: 4.71, Samples/sec: 0.48, Time/seq 2.10s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.92, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.27s, TFLOPs: 2.32, Samples/sec: 0.23, Time/seq 4.27s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.12s, TFLOPs: 4.67, Samples/sec: 0.47, Time/seq 2.12s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.92, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.88, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.37s, TFLOPs: 2.26, Samples/sec: 0.23, Time/seq 4.37s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.95, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.96s, TFLOPs: 5.04, Samples/sec: 0.51, Time/seq 1.96s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.94s, TFLOPs: 5.08, Samples/sec: 0.51, Time/seq 1.94s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.94s, TFLOPs: 5.09, Samples/sec: 0.52, Time/seq 1.94s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.29s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.29s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.17s, TFLOPs: 4.55, Samples/sec: 0.46, Time/seq 2.17s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.92, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.88, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.30s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.30s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.73, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.10s, TFLOPs: 4.71, Samples/sec: 0.48, Time/seq 2.10s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.54s, TFLOPs: 2.17, Samples/sec: 0.22, Time/seq 4.54s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.74, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.73, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.34s, TFLOPs: 4.22, Samples/sec: 0.43, Time/seq 2.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.12s, TFLOPs: 4.66, Samples/sec: 0.47, Time/seq 2.12s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.17s, TFLOPs: 2.37, Samples/sec: 0.24, Time/seq 4.17s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.93, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 4.99, Samples/sec: 0.50, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.03, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.03, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.88, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.87, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.30s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.30s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.74, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.36s, TFLOPs: 2.27, Samples/sec: 0.23, Time/seq 4.36s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.73, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.88, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.26s, TFLOPs: 2.32, Samples/sec: 0.23, Time/seq 4.26s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.31s, TFLOPs: 2.29, Samples/sec: 0.23, Time/seq 4.31s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.15s, TFLOPs: 4.60, Samples/sec: 0.46, Time/seq 2.15s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.52s, TFLOPs: 3.92, Samples/sec: 0.40, Time/seq 2.52s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.34s, TFLOPs: 2.28, Samples/sec: 0.23, Time/seq 4.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.79, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.93, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.28s, TFLOPs: 2.31, Samples/sec: 0.23, Time/seq 4.28s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.79, Samples/sec: 0.48, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.33s, TFLOPs: 2.28, Samples/sec: 0.23, Time/seq 4.33s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.77, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.87, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.25s, TFLOPs: 2.32, Samples/sec: 0.24, Time/seq 4.25s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.79, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.87, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.10s, TFLOPs: 4.71, Samples/sec: 0.48, Time/seq 2.10s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.35s, TFLOPs: 2.27, Samples/sec: 0.23, Time/seq 4.35s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.14s, TFLOPs: 4.61, Samples/sec: 0.47, Time/seq 2.14s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.77, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+[2026-01-29 22:33:59,925] [INFO] [launch.py:335:sigkill_handler] Killing subprocess 31271
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py", line 434, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py", line 387, in main
+[rank0]:     model.step()
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2690, in step
+[rank0]:     self._take_model_step(lr_kwargs)
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2585, in _take_model_step
+[rank0]:     self.optimizer.step()
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
+[rank0]:     ret_val = func(*args, **kwargs)
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py", line 2220, in step
+[rank0]:     self._reassign_or_swap_out_partitioned_parameters(sub_group_id)
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
+[rank0]:     ret_val = func(*args, **kwargs)
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py", line 2168, in _reassign_or_swap_out_partitioned_parameters
+[rank0]:     self.fp16_partitioned_groups_flat[sub_group_id].data.copy_(
+[rank0]: KeyboardInterrupt
+Traceback (most recent call last):
+  File "/home/ubuntu/.local/bin/deepspeed", line 6, in <module>
+    main()
+  File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/launcher/runner.py", line 646, in main
+    result.wait()
+  File "/usr/lib/python3.10/subprocess.py", line 1209, in wait
+    return self._wait(timeout=timeout)
+  File "/usr/lib/python3.10/subprocess.py", line 1959, in _wait
+    (pid, sts) = self._try_wait(0)
+  File "/usr/lib/python3.10/subprocess.py", line 1917, in _try_wait
+    (pid, sts) = os.waitpid(self.pid, wait_flags)
+KeyboardInterrupt
+[2026-01-29 22:34:00,546] [INFO] [launch.py:335:sigkill_handler] Killing subprocess 31271
+Exception ignored in atexit callback: <function shutdown_compile_workers at 0x7d22457a00d0>
+Traceback (most recent call last):
+  File "/usr/lib/python3/dist-packages/torch/_inductor/async_compile.py", line 113, in shutdown_compile_workers
+    pool.shutdown()
+  File "/usr/lib/python3/dist-packages/torch/_inductor/compile_worker/subproc_pool.py", line 239, in shutdown
+    self.process.wait(300)
+  File "/usr/lib/python3.10/subprocess.py", line 1209, in wait
+    return self._wait(timeout=timeout)
+  File "/usr/lib/python3.10/subprocess.py", line 1953, in _wait
+    time.sleep(delay)
+KeyboardInterrupt: 
+[2026-01-29 22:34:00,990] [INFO] [launch.py:335:sigkill_handler] Killing subprocess 31271
+[2026-01-29 22:34:04,967] [INFO] [launch.py:344:sigkill_handler] Main process received SIGINT, exiting

SFT-EN-01-29-2026/scripts/run_qwen3-4b.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Step 1: SFT Training for English Medical Data (UltraMedical)
+# Qwen3-4B with LoRA on H100
+
+MODEL_PATH=/workspace/Qwen3-4B
+DATA_PATH=/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl
+OUTPUT_DIR=./output_sft_en
+
+mkdir -p $OUTPUT_DIR
+
+deepspeed --num_gpus 1 main.py \
+   --model_name_or_path $MODEL_PATH \
+   --data_path $DATA_PATH \
+   --per_device_train_batch_size 2 \
+   --per_device_eval_batch_size 2 \
+   --max_seq_len 512 \
+   --learning_rate 2e-5 \
+   --weight_decay 0.1 \
+   --num_train_epochs 1 \
+   --num_warmup_steps 100 \
+   --gradient_accumulation_steps 4 \
+   --lr_scheduler_type cosine \
+   --gradient_checkpointing \
+   --dropout 0.0 \
+   --zero_stage 2 \
+   --dtype bf16 \
+   --lora_dim 64 \
+   --lora_module_name "layers." \
+   --only_optimize_lora \
+   --lora_learning_rate 5e-4 \
+   --compute_fp32_loss \
+   --print_loss \
+   --enable_tensorboard \
+   --tensorboard_path $OUTPUT_DIR \
+   --deepspeed \
+   --output_dir $OUTPUT_DIR

sft_model_backup/chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

sft_model_backup/config.json

@@ -0,0 +1,72 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "end_token_id": 151645,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9728,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 36,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "pad_token_id": 151645,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.0.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151672
+}

sft_model_backup/ds_tensorboard_logs/step1_model_tensorboard/events.out.tfevents.1769725308.209-20-158-64.30075.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bec6d2eb21f2a317beea31714d84e4c3fe34ba0c62365be1e2dc9ea98806cd55
+size 204

sft_model_backup/ds_tensorboard_logs/step1_model_tensorboard/events.out.tfevents.1769725536.209-20-158-64.31271.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:543d7327217546fd37d3fee5322ec0f7ccbb530b17e65a16f14250a52daa3a4a
+size 1448

sft_model_backup/ds_tensorboard_logs/step1_model_tensorboard/events.out.tfevents.1769726189.209-20-158-64.32221.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d996a4ea809b4cd08593bf3563f579d16402dfa2a1f23aa03d3110beda00a0e
+size 37198

sft_model_backup/ds_tensorboard_logs/step1_model_tensorboard/events.out.tfevents.1769727296.209-20-158-64.32989.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17e264bbafe4a0f20c4f2b0df948cee4cc696f181bd12f2530404e8c71c06444
+size 37198

sft_model_backup/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9956493572a7ae7ff86699c23789cba8d31a38d0a2d6333177d846b9d9cade23
+size 8820191160

sft_model_backup/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650

sft_model_backup/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "fast_tokenizer": true,
+  "is_local": true,
+  "model_max_length": 131072,
+  "pad_token": "<|im_end|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

sft_model_backup/training.log

@@ -0,0 +1,317 @@
+/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4
+  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"
+[2026-01-29 22:24:38,868] [WARNING] [runner.py:232:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.
+[2026-01-29 22:24:38,868] [INFO] [runner.py:630:main] cmd = /usr/bin/python3 -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbMF19 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None --log_level=info main.py --model_name_or_path /workspace/Qwen3-4B --data_path /home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl --weight_decay 0.1 --dropout 0.0 --gradient_accumulation_steps 8 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --zero_stage 3 --offload --dtype bf16 --enable_tensorboard --tensorboard_path ./output_sft_en --deepspeed --output_dir ./output_sft_en
+/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4
+  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"
+[2026-01-29 22:24:45,395] [INFO] [launch.py:162:main] WORLD INFO DICT: {'localhost': [0]}
+[2026-01-29 22:24:45,396] [INFO] [launch.py:168:main] nnodes=1, num_local_procs=1, node_rank=0
+[2026-01-29 22:24:45,396] [INFO] [launch.py:179:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0]})
+[2026-01-29 22:24:45,396] [INFO] [launch.py:180:main] dist_world_size=1
+[2026-01-29 22:24:45,396] [INFO] [launch.py:184:main] Setting CUDA_VISIBLE_DEVICES=0
+[2026-01-29 22:24:45,398] [INFO] [launch.py:272:main] process 31271 spawned with command: ['/usr/bin/python3', '-u', 'main.py', '--local_rank=0', '--model_name_or_path', '/workspace/Qwen3-4B', '--data_path', '/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl', '--weight_decay', '0.1', '--dropout', '0.0', '--gradient_accumulation_steps', '8', '--per_device_train_batch_size', '1', '--per_device_eval_batch_size', '1', '--zero_stage', '3', '--offload', '--dtype', 'bf16', '--enable_tensorboard', '--tensorboard_path', './output_sft_en', '--deepspeed', '--output_dir', './output_sft_en']
+/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4
+  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"
+[rank0]:[W129 22:24:52.444107661 ProcessGroupNCCL.cpp:4715] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 as device used by this process is currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. You can pecify device_id in init_process_group() to force use of a particular device.
+Setting model_config.attention_dropout to 0.0
+args:  Namespace(data_path=['/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl'], data_split='6,2,2', sft_only_data_path=[], data_output_path='/tmp/data_files/', model_name_or_path='/workspace/Qwen3-4B', per_device_train_batch_size=1, per_device_eval_batch_size=1, max_seq_len=512, learning_rate=0.001, weight_decay=0.1, num_train_epochs=1, gradient_accumulation_steps=8, lr_scheduler_type=<SchedulerType.COSINE: 'cosine'>, num_warmup_steps=0, output_dir='./output_sft_en', seed=1234, local_rank=0, gradient_checkpointing=False, dropout=0.0, offload=True, dtype='bf16', zero_stage=3, lora_dim=0, lora_module_name='decoder.layers.', only_optimize_lora=False, lora_learning_rate=0.0005, compute_fp32_loss=False, enable_tensorboard=True, tensorboard_path='./output_sft_en', add_eot_token=False, eot_token='<|endoftext|>', print_loss=False, deepspeed=True, deepspeed_config=None, deepscale=False, deepscale_config=None, global_rank=0)
+data_path:  ['/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/data/train.jsonl']
+/usr/lib/python3/dist-packages/torch/utils/cpp_extension.py:2376: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. 
+If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].
+  warnings.warn(
+2026-01-29 22:25:34.798274: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+2026-01-29 22:25:34.808869: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1769725534.821805   31271 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1769725534.825823   31271 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1769725534.835606   31271 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1769725534.835626   31271 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1769725534.835656   31271 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1769725534.835658   31271 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+2026-01-29 22:25:34.838493: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
+To enable the following instructions: AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI, in other operations, rebuild TensorFlow with the appropriate compiler flags.
+Stage 3 initialize beginning
+MA 0.72 GB         Max_MA 2.9 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 16.26 GB, percent = 7.4%
+DeepSpeedZeRoOffload initialize [begin]
+MA 0.72 GB         Max_MA 0.72 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 16.25 GB, percent = 7.3%
+Parameter Offload - Persistent parameters statistics: param_count = 145, numel = 196096
+DeepSpeedZeRoOffload initialize [end]
+MA 0.0 GB         Max_MA 0.72 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 16.7 GB, percent = 7.6%
+Before creating fp16 partitions
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 16.7 GB, percent = 7.6%
+After creating fp16 partitions: 5
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 19.89 GB, percent = 9.0%
+Before creating fp32 partitions
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 19.89 GB, percent = 9.0%
+After creating fp32 partitions
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 34.0 GB, percent = 15.4%
+Before initializing optimizer states
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 34.0 GB, percent = 15.4%
+After initializing optimizer states
+MA 0.0 GB         Max_MA 0.0 GB         CA 2.9 GB         Max_CA 3 GB 
+CPU Virtual Memory:  used = 49.09 GB, percent = 22.2%
+After initializing ZeRO optimizer
+MA 0.93 GB         Max_MA 2.38 GB         CA 3.83 GB         Max_CA 4 GB 
+CPU Virtual Memory:  used = 56.32 GB, percent = 25.5%
+***** Running training *****
+Beginning of Epoch 1/1, Total Micro Batches 5400
+Model Parameters: 4.022 B, Latency: 2.91s, TFLOPs: 3.40, Samples/sec: 0.34, Time/seq 2.91s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 3.07s, TFLOPs: 3.22, Samples/sec: 0.33, Time/seq 3.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.34s, TFLOPs: 4.22, Samples/sec: 0.43, Time/seq 2.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.35s, TFLOPs: 4.20, Samples/sec: 0.43, Time/seq 2.35s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.34s, TFLOPs: 4.23, Samples/sec: 0.43, Time/seq 2.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.34s, TFLOPs: 4.22, Samples/sec: 0.43, Time/seq 2.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.33s, TFLOPs: 4.23, Samples/sec: 0.43, Time/seq 2.33s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 6.18s, TFLOPs: 1.60, Samples/sec: 0.16, Time/seq 6.18s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.11s, TFLOPs: 4.69, Samples/sec: 0.47, Time/seq 2.11s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 4.99, Samples/sec: 0.50, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.95, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.96s, TFLOPs: 5.04, Samples/sec: 0.51, Time/seq 1.96s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 4.99, Samples/sec: 0.50, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.19s, TFLOPs: 2.36, Samples/sec: 0.24, Time/seq 4.19s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 5.00, Samples/sec: 0.51, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.97, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.97, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.02, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.92, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.21s, TFLOPs: 2.35, Samples/sec: 0.24, Time/seq 4.21s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.02, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.93, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.02, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 5.00, Samples/sec: 0.51, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.97, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.01, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.24s, TFLOPs: 2.33, Samples/sec: 0.24, Time/seq 4.24s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.39s, TFLOPs: 4.14, Samples/sec: 0.42, Time/seq 2.39s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.36s, TFLOPs: 4.19, Samples/sec: 0.42, Time/seq 2.36s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.31s, TFLOPs: 4.27, Samples/sec: 0.43, Time/seq 2.31s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.96, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.93, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.27s, TFLOPs: 2.31, Samples/sec: 0.23, Time/seq 4.27s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.95s, TFLOPs: 5.06, Samples/sec: 0.51, Time/seq 1.95s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.94s, TFLOPs: 5.09, Samples/sec: 0.52, Time/seq 1.94s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.93s, TFLOPs: 5.12, Samples/sec: 0.52, Time/seq 1.93s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.97, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.27s, TFLOPs: 2.31, Samples/sec: 0.23, Time/seq 4.27s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.13s, TFLOPs: 4.64, Samples/sec: 0.47, Time/seq 2.13s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 4.98, Samples/sec: 0.50, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.99s, TFLOPs: 4.98, Samples/sec: 0.50, Time/seq 1.99s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.74, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.22s, TFLOPs: 2.34, Samples/sec: 0.24, Time/seq 4.22s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.74, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.75, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.77, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.10s, TFLOPs: 4.71, Samples/sec: 0.48, Time/seq 2.10s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.75, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.30s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.30s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.77, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.25s, TFLOPs: 4.40, Samples/sec: 0.45, Time/seq 2.25s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.29s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.29s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.39s, TFLOPs: 4.13, Samples/sec: 0.42, Time/seq 2.39s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.37s, TFLOPs: 4.17, Samples/sec: 0.42, Time/seq 2.37s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.37s, TFLOPs: 4.18, Samples/sec: 0.42, Time/seq 2.37s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.37s, TFLOPs: 4.17, Samples/sec: 0.42, Time/seq 2.37s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.28s, TFLOPs: 4.33, Samples/sec: 0.44, Time/seq 2.28s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.87, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.28s, TFLOPs: 2.31, Samples/sec: 0.23, Time/seq 4.28s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.10s, TFLOPs: 4.71, Samples/sec: 0.48, Time/seq 2.10s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.92, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.27s, TFLOPs: 2.32, Samples/sec: 0.23, Time/seq 4.27s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.12s, TFLOPs: 4.67, Samples/sec: 0.47, Time/seq 2.12s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.92, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.88, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.37s, TFLOPs: 2.26, Samples/sec: 0.23, Time/seq 4.37s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.95, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.96s, TFLOPs: 5.04, Samples/sec: 0.51, Time/seq 1.96s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.94s, TFLOPs: 5.08, Samples/sec: 0.51, Time/seq 1.94s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.94s, TFLOPs: 5.09, Samples/sec: 0.52, Time/seq 1.94s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.29s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.29s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.17s, TFLOPs: 4.55, Samples/sec: 0.46, Time/seq 2.17s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.92, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.88, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.30s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.30s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.73, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.10s, TFLOPs: 4.71, Samples/sec: 0.48, Time/seq 2.10s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.54s, TFLOPs: 2.17, Samples/sec: 0.22, Time/seq 4.54s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.74, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.73, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.34s, TFLOPs: 4.22, Samples/sec: 0.43, Time/seq 2.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.12s, TFLOPs: 4.66, Samples/sec: 0.47, Time/seq 2.12s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.17s, TFLOPs: 2.37, Samples/sec: 0.24, Time/seq 4.17s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.93, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.98s, TFLOPs: 4.99, Samples/sec: 0.50, Time/seq 1.98s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.03, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 1.97s, TFLOPs: 5.03, Samples/sec: 0.51, Time/seq 1.97s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.88, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.87, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.30s, TFLOPs: 2.30, Samples/sec: 0.23, Time/seq 4.30s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.74, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.36s, TFLOPs: 2.27, Samples/sec: 0.23, Time/seq 4.36s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.09s, TFLOPs: 4.73, Samples/sec: 0.48, Time/seq 2.09s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.88, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.26s, TFLOPs: 2.32, Samples/sec: 0.23, Time/seq 4.26s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.91, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.31s, TFLOPs: 2.29, Samples/sec: 0.23, Time/seq 4.31s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.15s, TFLOPs: 4.60, Samples/sec: 0.46, Time/seq 2.15s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.52s, TFLOPs: 3.92, Samples/sec: 0.40, Time/seq 2.52s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.34s, TFLOPs: 2.28, Samples/sec: 0.23, Time/seq 4.34s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.00s, TFLOPs: 4.94, Samples/sec: 0.50, Time/seq 2.00s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.79, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.01s, TFLOPs: 4.93, Samples/sec: 0.50, Time/seq 2.01s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.28s, TFLOPs: 2.31, Samples/sec: 0.23, Time/seq 4.28s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.79, Samples/sec: 0.48, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.89, Samples/sec: 0.49, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.78, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.33s, TFLOPs: 2.28, Samples/sec: 0.23, Time/seq 4.33s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.77, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.87, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.02s, TFLOPs: 4.90, Samples/sec: 0.50, Time/seq 2.02s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.82, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.25s, TFLOPs: 2.32, Samples/sec: 0.24, Time/seq 4.25s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.83, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.79, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.84, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.87, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.10s, TFLOPs: 4.71, Samples/sec: 0.48, Time/seq 2.10s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.03s, TFLOPs: 4.86, Samples/sec: 0.49, Time/seq 2.03s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 4.35s, TFLOPs: 2.27, Samples/sec: 0.23, Time/seq 4.35s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.14s, TFLOPs: 4.61, Samples/sec: 0.47, Time/seq 2.14s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.08s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.08s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.06s, TFLOPs: 4.80, Samples/sec: 0.49, Time/seq 2.06s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.76, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.04s, TFLOPs: 4.85, Samples/sec: 0.49, Time/seq 2.04s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.07s, TFLOPs: 4.77, Samples/sec: 0.48, Time/seq 2.07s, Batch Size: 1, Sequence Length: 512
+Model Parameters: 4.022 B, Latency: 2.05s, TFLOPs: 4.81, Samples/sec: 0.49, Time/seq 2.05s, Batch Size: 1, Sequence Length: 512
+[2026-01-29 22:33:59,925] [INFO] [launch.py:335:sigkill_handler] Killing subprocess 31271
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py", line 434, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py", line 387, in main
+[rank0]:     model.step()
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2690, in step
+[rank0]:     self._take_model_step(lr_kwargs)
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2585, in _take_model_step
+[rank0]:     self.optimizer.step()
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
+[rank0]:     ret_val = func(*args, **kwargs)
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py", line 2220, in step
+[rank0]:     self._reassign_or_swap_out_partitioned_parameters(sub_group_id)
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
+[rank0]:     ret_val = func(*args, **kwargs)
+[rank0]:   File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py", line 2168, in _reassign_or_swap_out_partitioned_parameters
+[rank0]:     self.fp16_partitioned_groups_flat[sub_group_id].data.copy_(
+[rank0]: KeyboardInterrupt
+Traceback (most recent call last):
+  File "/home/ubuntu/.local/bin/deepspeed", line 6, in <module>
+    main()
+  File "/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/launcher/runner.py", line 646, in main
+    result.wait()
+  File "/usr/lib/python3.10/subprocess.py", line 1209, in wait
+    return self._wait(timeout=timeout)
+  File "/usr/lib/python3.10/subprocess.py", line 1959, in _wait
+    (pid, sts) = self._try_wait(0)
+  File "/usr/lib/python3.10/subprocess.py", line 1917, in _try_wait
+    (pid, sts) = os.waitpid(self.pid, wait_flags)
+KeyboardInterrupt
+[2026-01-29 22:34:00,546] [INFO] [launch.py:335:sigkill_handler] Killing subprocess 31271
+Exception ignored in atexit callback: <function shutdown_compile_workers at 0x7d22457a00d0>
+Traceback (most recent call last):
+  File "/usr/lib/python3/dist-packages/torch/_inductor/async_compile.py", line 113, in shutdown_compile_workers
+    pool.shutdown()
+  File "/usr/lib/python3/dist-packages/torch/_inductor/compile_worker/subproc_pool.py", line 239, in shutdown
+    self.process.wait(300)
+  File "/usr/lib/python3.10/subprocess.py", line 1209, in wait
+    return self._wait(timeout=timeout)
+  File "/usr/lib/python3.10/subprocess.py", line 1953, in _wait
+    time.sleep(delay)
+KeyboardInterrupt: 
+[2026-01-29 22:34:00,990] [INFO] [launch.py:335:sigkill_handler] Killing subprocess 31271
+[2026-01-29 22:34:04,967] [INFO] [launch.py:344:sigkill_handler] Main process received SIGINT, exiting