add sft training code using loomtrain.core

iiGray · iiGray · commit 3f0ba85133a9 · 2025-11-02T15:58:09.000+08:00
diff --git a/examples/train_sft_core.sh b/examples/train_sft_core.sh
@@ -0,0 +1,17 @@
+TORCH_CUDA_ARCH_LIST="8.6;8.9" deepspeed --module loomtrain.scripts.train_sft \
+    --model-path /data/hf_models/Meta-Llama-3.1-8B-Instruct/ \
+    --dataset-paths /data/lcm_lab/jbb/datas/Skywork-Reward-80k-Formatted \
+    --prompt-key chat_template \
+    --response-key chosen \
+    --train-samples  800\
+    --val-samples 80 \
+    --max-length 128000 \
+    --packing-length 0 \
+    --micro-batch-size 1 \
+    --global-batch-size 16 \
+    --val-batch-size 8 \
+    --cp-size 1 \
+    --val-interval 20 \
+    --ckpt-interval 20 \
+    --weight-interval 20 \
+    --save-dir ./Test-Loom-Llama3.1 \
diff --git a/loomtrain/core/tasks/sft.py b/loomtrain/core/tasks/sft.py
@@ -24,13 +24,13 @@ def setup_self_module(self):
         self.actor = self.opt_groups['group0'].actor
         self.toknizer = self.opt_groups['group0'].tokenizer
         self.optimizer = self.opt_groups['group0'].optimizer
-        self.scheduler = self.opt_groups['group0'].scheduler, 
+        self.scheduler = self.opt_groups['group0'].scheduler
         self.loss_fn = self.opt_groups['group0'].loss_fn
 
     def micro_batch_forward_backward(self, batch) -> "dict[str, object]":
-        inputs, attention_masks, loss_masks, seq_lens = batch
-        output = self.actor(sequences = inputs, attention_masks = attention_masks,seq_lens = seq_lens)
-        labels = torch.where(attention_masks.bool() & loss_masks.bool(), inputs, self.loss_fn.ignore_index)
+        inputs, attention_mask, loss_mask, seq_lens = batch
+        output = self.actor(sequences = inputs, attention_mask = attention_mask, seq_lens = seq_lens)
+        labels = torch.where(attention_mask.bool() & loss_mask.bool(), inputs, self.loss_fn.ignore_index)
 
         gpt_loss = self.loss_fn(output.logits, labels)
 
@@ -39,7 +39,7 @@ def micro_batch_forward_backward(self, batch) -> "dict[str, object]":
         return dict(
             loss = gpt_loss.item(),
             total_tokens = parallel.all_reduce(sum(seq_lens)) * parallel.get_dp_count() / 10 ** 9,
-            loss_tokens = parallel.all_reduce(loss_masks.int().sum().item()) * parallel.get_dp_count() / 10 ** 9
+            loss_tokens = parallel.all_reduce(loss_mask.int().sum().item()) * parallel.get_dp_count() / 10 ** 9
         )
 
     def micro_batch_validate_forward(self, batch):
@@ -74,55 +74,6 @@ def __init__(self,
             assert "prompt_key" in data_dict
             assert  "response_key" in data_dict
 
-    @LoomDataModule.datasetmethod
-    def filter_data(dataset, self:"LoomSFTData", data: "dict"):
-        if dataset.max_length < 128000:
-            prompt_template = data[dataset.prompt_key]
-            response_template = role_template(data[dataset.response_key], "assistant")
-            tokenized = dataset.tokenizer.apply_chat_template(
-                prompt_template + response_template, tokenize = True, 
-                max_length = 128000, padding = False,
-                truncation = True, return_tensors = 'pt'
-            )
-            if tokenized.numel() > dataset.max_length: return False
-        return True
-
-    @LoomDataModule.datasetmethod
-    def process_data(dataset, data):
-        prompt_template = data[dataset.prompt_key]
-        response_text = data[dataset.response_key]
-        if isinstance(response_text, str):
-            response_text = [{"role":"assistant", "content": response_text}]
-        prompt = dataset.tokenizer.apply_chat_template(
-            prompt_template, tokenize = False, add_generation_prompt = True
-        )
-        response = dataset.tokenizer.apply_chat_template(
-            prompt_template + response_text, tokenize = False
-        )[len(prompt): ]
-
-
-        prompt_token = dataset.tokenizer(prompt, max_length = dataset.max_length,
-                                      padding = False,
-                                      truncation = True,
-                                      return_tensors = 'pt',
-                                      add_special_tokens = False)
-        response_token = dataset.tokenizer(response, max_length = dataset.max_length,
-                                        padding = False,
-                                        truncation = True,
-                                        return_tensors = 'pt',
-                                        add_special_tokens = False)
-        
-        prompt_ids_len = prompt_token["attention_mask"].int().sum().item()
-        input_ids_len = prompt_ids_len + response_token["attention_mask"].int().sum().item()
-
-        return dict(
-            prompt = prompt,
-            response = response,
-            prompt_ids_len = prompt_ids_len,
-            input_ids_len = input_ids_len,
-            response_ranges = None # not multiturn
-        )
-
     @LoomDataModule.datasetmethod
     def get_loss_mask(dataset, input_ids, idx):
         loss_mask = torch.zeros_like(input_ids, dtype = torch.bool)
@@ -132,19 +83,70 @@ def get_loss_mask(dataset, input_ids, idx):
 
         return loss_mask
 
+    def dataset_initialize(dataset, self: "LoomSFTData", raw_dataset, data_dict):
 
+        tokenizer = dataset.tokenizer = data_dict.tokenizer
+        prompt_key = dataset.prompt_key = data_dict["prompt_key"]
+        response_key = dataset.response_key = data_dict["response_key"]
+        max_length = dataset.max_length = self.max_length
+
+        def filter_data(data: "dict"):
+            if max_length < 128000:
+                prompt_template = data[prompt_key]
+                response_template = role_template(data[response_key], "assistant")
+                tokenized = tokenizer.apply_chat_template(
+                    prompt_template + response_template, tokenize = True, 
+                    max_length = 128000, padding = False,
+                    truncation = True, return_tensors = 'pt'
+                )
+                if tokenized.numel() > max_length: return False
+            return True
+
+        def process_data(data):
+            prompt_template = data[prompt_key]
+            response_text = data[response_key]
+            if isinstance(response_text, str):
+                response_text = [{"role":"assistant", "content": response_text}]
+            prompt = tokenizer.apply_chat_template(
+                prompt_template, tokenize = False, add_generation_prompt = True
+            )
+            response = tokenizer.apply_chat_template(
+                prompt_template + response_text, tokenize = False
+            )[len(prompt): ]
 
-    def dataset_initialize(dataset, self: "LoomSFTData", raw_dataset, data_dict):
-        dataset.tokenizer = data_dict.tokenizer
-        raw_dataset = raw_dataset.filter(dataset.filter_data, num_proc = self.num_proc)
-        processed_dataset = raw_dataset.map(dataset.process_data,
+
+            prompt_token = tokenizer(prompt, max_length = max_length,
+                                        padding = False,
+                                        truncation = True,
+                                        return_tensors = 'pt',
+                                        add_special_tokens = False)
+            response_token = tokenizer(response, max_length = max_length,
+                                            padding = False,
+                                            truncation = True,
+                                            return_tensors = 'pt',
+                                            add_special_tokens = False)
+            
+            prompt_ids_len = prompt_token["attention_mask"].int().sum().item()
+            input_ids_len = prompt_ids_len + response_token["attention_mask"].int().sum().item()
+
+            return dict(
+                prompt = prompt,
+                response = response,
+                prompt_ids_len = prompt_ids_len,
+                input_ids_len = input_ids_len,
+                response_ranges = None # not multiturn
+            )
+
+        raw_dataset = raw_dataset.filter(filter_data, num_proc = self.num_proc)
+        processed_dataset = raw_dataset.map(process_data,
                                             remove_columns = raw_dataset.column_names,
                                             num_proc = self.num_proc)
         
         dataset.prompts = processed_dataset["prompt"]
         dataset.responses = processed_dataset["response"]
         dataset.prompt_ids_lens = processed_dataset["prompt_ids_len"]
-        dataset.input_ids_lens = processed_dataset["prompt_ids_len"]
+        # This attribute makes itself possible to be packed.
+        dataset._input_ids_lens = processed_dataset["prompt_ids_len"]
 
 
     
@@ -193,7 +195,7 @@ def dataset_collate_fn(dataset, self: "LoomSFTData", item_list):
 
         if packed_input_ids.numel() % self.cp_size:
             padding_len = self.cp_size - (packed_input_ids.numel() % self.cp_size)
-            packed_input_ids = F.pad(packed_input_ids, (0, padding_len), value = self.tokenizer.pad_token_id)
+            packed_input_ids = F.pad(packed_input_ids, (0, padding_len), value = dataset.tokenizer.pad_token_id)
             packed_attention_masks = F.pad(packed_attention_masks, (0, padding_len), value = 0)
             packed_loss_masks = F.pad(packed_loss_masks, (0, padding_len), value = 0)
 
diff --git a/loomtrain/scripts/train_sft.py b/loomtrain/scripts/train_sft.py
@@ -1,3 +1,5 @@
+import argparse, json
+
 from loomtrain.core import (
     LoomSFTModule,
     LoomSFTData,
@@ -8,50 +10,132 @@
     SortPackingStrategy,
     DataConfig,
     CheckpointConfig,
+    VisualizationModule,
     parallel
 
 )
 
 
-if __name__ == "__main__":
-    args = ...
+def train(args):
+
+    data_config = DataConfig(
+        collate_type = 'packing',
+        packing_length = args.packing_length,
+        
+        train_batch_size = args.global_batch_size,
+        micro_batch_size = args.micro_batch_size,
+        val_batch_size = args.val_batch_size,
+        val_interval = args.val_interval,
+        batch_size = 10,
+        num_epochs = 1
+    )
+
 
     parallel_config = parallel.ParallelConfig(
-        train_batch_size = 2,
-        micro_batch_size = 1,
-        val_batch_size = 2,
         nnodes = 1,
         devices_per_node = 8,
-        cp = 8
+        cp = args.cp_size,
+        cp_type = args.cp_type,
+        cp_args = args.cp_args
     )
 
-    data_config = DataConfig(
-        collate_type = 'packing',
-        packing_length = 64000,
-        val_interval = 20,
-        batch_size = 10,
-        num_epochs = 1
+    checkpoint_config = CheckpointConfig(
+        save_dir = args.save_dir,
+        ckpt_interval = args.ckpt_interval,
+        weight_interval = args.weight_interval,
     )
 
-    train_strategy = DeepspeedStrategy(parallel_config, data_config = data_config)
+
+    train_strategy = DeepspeedStrategy(parallel_config, data_config = data_config,
+                                       deepspeed_config= DeepspeedConfig(zero_stage =3, adam_offload = True))
 
     data_strategy = SortPackingStrategy(parallel_config, data_config = data_config)
 
     trainer = LoomTrainer(train_strategy = train_strategy,
                           data_strategy = data_strategy)
     
-    module = LoomSFTModule("/path/to/model/")
+    module = LoomSFTModule(args.model_path)
     
     datamodule = LoomSFTData([
-        LoomDataDict(path = "/path/to/data",count = 10, prompt_key = 'prompt', response_key = 'response'),
-        ...
-    ], max_length = 128000)
+        LoomDataDict(data_path = pth, tokenizer_path = args.model_path, train_count = tc, val_count = vc, prompt_key = args.prompt_key, response_key = args.response_key) for pth, tc, vc in zip(
+            args.dataset_paths, args.train_samples, args.val_samples
+        )
+    ], max_length = args.max_length)
 
 
-    ckpt_cfg = CheckpointConfig(load_dir = '/path/to/load', save_dir = '/path/to/save')
+    vismodule = VisualizationModule(logtype = 'tensorboard')
+
 
     trainer.fit(
         module = module,
         datamodule = datamodule,
-        ckpt_cfg = ckpt_cfg
+        vismodule = vismodule,
+        checkpoint_config = checkpoint_config
+    )
+
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--local_rank", type = int , default = -1
     )
+    
+    parser.add_argument(
+        "--global-batch-size", type = int , default = 64
+    )
+    parser.add_argument(
+        "--micro-batch-size", type = int , default = 1
+    )
+    parser.add_argument(
+        "--val-batch-size", type = int, default = 1,
+    )
+    parser.add_argument(
+        "--cp-size", type = int, default = 8
+    )
+    parser.add_argument(
+        "--cp-type", type = str, default = "ring"
+    )
+    parser.add_argument(
+        "--cp-args", type = json.loads, default = '{"head_stride": 1}'
+    )
+    parser.add_argument(
+        "--max-length", type = int, default = 128000
+    )
+    parser.add_argument(
+        "--packing-length", type = int, default = None
+    )
+    parser.add_argument(
+        "--val-interval", type = int, default = 20,
+    )
+    parser.add_argument(
+        "--save-dir", type = str, required = True,
+    )
+    parser.add_argument(
+        "--ckpt-interval", type = int, default = 20,
+    )
+    parser.add_argument(
+        "--weight-interval", type = int, default = 20,
+    )
+    parser.add_argument(
+        "--model-path", type = str, required = True
+    )
+    parser.add_argument(
+        "--dataset-paths", type = str, nargs = "+", required = True
+    )
+    parser.add_argument(
+        "--train-samples", type = int, nargs = "+", required = True
+    )
+    parser.add_argument(
+        "--val-samples", type = int, nargs = "+", required = True
+    )
+    parser.add_argument(
+        "--prompt-key", type = str, default = "prompt"
+    )
+    parser.add_argument(
+        "--response-key", type = str, default = "response"
+    )
+
+    args = parser.parse_args()
+    train(args)
