diff --git a/1-pretrain.py b/1-pretrain.py index d6bfff6..ad1a45c 100644 --- a/1-pretrain.py +++ b/1-pretrain.py @@ -74,7 +74,8 @@ def train_epoch(epoch, wandb): loss.item() * args.accumulation_steps, optimizer.param_groups[-1]['lr'], spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60)) - if wandb is not None: + + if (wandb is not None) and (not ddp or dist.get_rank() == 0): wandb.log({"loss": loss.item() * args.accumulation_steps, "lr": optimizer.param_groups[-1]['lr'], "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60}) @@ -115,6 +116,7 @@ def init_distributed_mode(): DEVICE = f"cuda:{ddp_local_rank}" torch.cuda.set_device(DEVICE) + # torchrun --nproc_per_node 2 1-pretrain.py if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind Pretraining") @@ -149,6 +151,7 @@ if __name__ == "__main__": args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}" ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast() + ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run? ddp_local_rank, DEVICE = 0, "cuda:0" if ddp: diff --git a/3-full_sft.py b/3-full_sft.py index fc58bb3..d82e662 100644 --- a/3-full_sft.py +++ b/3-full_sft.py @@ -80,8 +80,9 @@ def train_epoch(epoch, wandb): loss.item(), optimizer.param_groups[-1]['lr'], spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60)) - if wandb is not None: - wandb.log({"loss": loss.item(), + + if (wandb is not None) and (not ddp or dist.get_rank() == 0): + wandb.log({"loss": loss, "lr": optimizer.param_groups[-1]['lr'], "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60}) diff --git a/4-lora_sft.py b/4-lora_sft.py index 44d69ba..e72f8ca 100644 --- a/4-lora_sft.py +++ b/4-lora_sft.py @@ -102,8 +102,8 @@ def find_all_linear_names(model): def init_model(): - model_name_or_path = "./minimind" - tokenizer_name_or_path = "./minimind" + model_name_or_path = "./minimind-v1-small" + tokenizer_name_or_path = "./minimind-v1-small" tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True, use_fast=False) model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True).to(args.device) diff --git a/README.md b/README.md index 92ebcfc..265b5a0 100644 --- a/README.md +++ b/README.md @@ -69,10 +69,9 @@ https://github.com/user-attachments/assets/88b98128-636e-43bc-a419-b1b1403c2055 - 公开MiniMind模型代码(包含Dense和MoE模型)、Pretrain、SFT指令微调、LoRA微调、DPO偏好优化的全过程代码、数据集和来源。 - 兼容`transformers`、`accelerate`、`trl`、`peft`等流行框架。 -- 训练支持单机单卡、单机多卡(DDP、DeepSpeed)训练。训练过程中支持在任意位置停止,及在任意位置继续训练。 +- 训练支持单机单卡、单机多卡(DDP、DeepSpeed)训练,使用wandb可视化训练流程。支持在任意位置停止,及在任意位置继续训练。 - 在Ceval数据集上进行模型测试的代码。 - 实现Openai-Api基本的chat接口,便于集成到第三方ChatUI使用(FastGPT、Open-WebUI等)。 -- 使用wandb可视化训练流程。 希望此开源项目可以帮助LLM初学者快速入门! @@ -696,6 +695,8 @@ minimind模型本身没有使用较大的数据集训练,也没有针对回答     + +  ## 😊鸣谢 diff --git a/README_en.md b/README_en.md index af02bc7..a0fa5a4 100644 --- a/README_en.md +++ b/README_en.md @@ -75,13 +75,10 @@ The project includes: - Public MiniMind model code (including Dense and MoE models), code for Pretrain, SFT instruction fine-tuning, LoRA fine-tuning, and DPO preference optimization, along with datasets and sources. - Compatibility with popular frameworks such as `transformers`, `accelerate`, `trl`, and `peft`. -- Training support for single-GPU and multi-GPU setups(DDP、DeepSpeed). The training process allows for stopping and - resuming at any - point. +- Training support for single-GPU and multi-GPU setups(DDP、DeepSpeed), Use wandb to visualize the training process. The training process allows for stopping and resuming at any point. - Code for testing the model on the Ceval dataset. - Implementation of a basic chat interface compatible with OpenAI's API, facilitating integration into third-party Chat UIs (such as FastGPT, Open-WebUI, etc.). -- Use wandb to visualize the training process. We hope this open-source project helps LLM beginners get started quickly! @@ -764,6 +761,8 @@ your model with third-party UIs, such as fastgpt, OpenWebUI, etc.     + +  ## 😊Thanks for diff --git a/data_process.py b/data_process.py index 9c03628..047ff0e 100644 --- a/data_process.py +++ b/data_process.py @@ -95,7 +95,7 @@ def process_seq_monkey(chunk_size=50000): if doc_ids: arr = np.array(doc_ids, dtype=np.uint16) - with open(f'./dataset/clean_seq_monkey.bin', 'wb') as f: + with open(f'./dataset/clean_seq_monkey.bin', 'ab') as f: f.write(arr.tobytes())