Merge 66aef88185 into f659b55761

[fix] save final checkpoints after the residual accumulation step

README.md

@@ -27,7 +27,7 @@
 
 <div align="center">
 
-中文 | [English](./README_en.md)
+中文 | [English](./README_en.md) | [العربية](./README_i18n.md#arabic) | [Français](./README_i18n.md#french) | [Español](./README_i18n.md#spanish) | [Português](./README_i18n.md#portuguese) | [Deutsch](./README_i18n.md#german) | [فارسی](./README_i18n.md#persian) | [日本語](./README_i18n.md#japanese) | [한국어](./README_i18n.md#korean)
 
 </div>
 

README_en.md

@@ -27,7 +27,7 @@
 
 <div align="center">
 
-[中文](./README.md) | English
+[中文](./README.md) | English | [العربية](./README_i18n.md#arabic) | [Français](./README_i18n.md#french) | [Español](./README_i18n.md#spanish) | [Português](./README_i18n.md#portuguese) | [Deutsch](./README_i18n.md#german) | [فارسی](./README_i18n.md#persian) | [日本語](./README_i18n.md#japanese) | [한국어](./README_i18n.md#korean)
 
 </div>
 

README_i18n.md

@@ -0,0 +1,197 @@
+# MiniMind in other languages
+
+[中文](./README.md) | [English](./README_en.md) | [العربية](#arabic) | [Français](#french) | [Español](#spanish) | [Português](#portuguese) | [Deutsch](#german) | [فارسی](#persian) | [日本語](#japanese) | [한국어](#korean)
+
+These localized introductions provide a concise overview and the fastest path to running MiniMind. The complete and most current technical documentation is maintained in [English](./README_en.md) and [Chinese](./README.md).
+
+<a id="arabic"></a>
+
+<h2 dir="rtl">العربية</h2>
+
+<p dir="rtl">MiniMind هو مشروع مفتوح المصدر يتيح تدريب نموذج لغوي صغير يضم نحو 64 مليون مُعلَمة من الصفر. يمكن إكمال مرحلة الضبط الدقيق الخاضع للإشراف خلال ساعتين تقريبًا على بطاقة NVIDIA 3090 واحدة، بتكلفة استئجار تقارب 3 يوانات صينية.</p>
+
+<ul dir="rtl">
+  <li>تنفيذ مبسّط ومباشر باستخدام PyTorch لبنية Dense وMoE.</li>
+  <li>مسار تدريب متكامل يشمل التدريب المسبق وSFT وLoRA وDPO وPPO وGRPO وCISPO واستخدام الأدوات وAgentic RL والتقطير.</li>
+  <li>دعم Transformers وllama.cpp وvLLM وOllama وواجهة API متوافقة مع OpenAI وواجهة WebUI.</li>
+  <li>بيانات وأوزان ونصوص تدريب مفتوحة لتسهيل التعلم وإعادة إنتاج النتائج.</li>
+</ul>
+
+<p dir="rtl"><strong>البدء السريع</strong></p>
+
+```bash
+git clone --depth 1 https://github.com/jingyaogong/minimind
+cd minimind
+pip install -r requirements.txt
+modelscope download --model gongjy/minimind-3 --local_dir ./minimind-3
+python eval_llm.py --load_from ./minimind-3
+```
+
+<p dir="rtl">للحصول على تعليمات التدريب ومجموعات البيانات والتقييم وتحويل النموذج، راجع <a href="./README_en.md">الدليل الإنجليزي الكامل</a>.</p>
+
+<a id="french"></a>
+
+## Français
+
+MiniMind est un projet open source permettant d'entraîner de zéro un petit modèle de langage d'environ 64 millions de paramètres. L'étape de fine-tuning supervisé peut être réalisée en environ deux heures sur une seule NVIDIA 3090, pour un coût de location proche de 3 RMB.
+
+- Implémentation PyTorch minimale et lisible des architectures Dense et MoE.
+- Pipeline complet : préentraînement, SFT, LoRA, DPO, PPO, GRPO, CISPO, Tool Use, Agentic RL et distillation.
+- Compatibilité avec Transformers, llama.cpp, vLLM et Ollama, ainsi qu'une API compatible OpenAI et une WebUI.
+- Jeux de données, poids et scripts d'entraînement ouverts pour apprendre et reproduire les résultats.
+
+**Démarrage rapide**
+
+```bash
+git clone --depth 1 https://github.com/jingyaogong/minimind
+cd minimind
+pip install -r requirements.txt
+modelscope download --model gongjy/minimind-3 --local_dir ./minimind-3
+python eval_llm.py --load_from ./minimind-3
+```
+
+Pour les instructions complètes sur l'entraînement, les données, l'évaluation et la conversion des modèles, consultez le [guide complet en anglais](./README_en.md).
+
+<a id="spanish"></a>
+
+## Español
+
+MiniMind es un proyecto de código abierto para entrenar desde cero un modelo de lenguaje pequeño de unos 64 millones de parámetros. La etapa de ajuste supervisado puede completarse en unas dos horas con una sola NVIDIA 3090, con un coste de alquiler cercano a 3 RMB.
+
+- Implementación mínima y legible en PyTorch de arquitecturas Dense y MoE.
+- Flujo completo: preentrenamiento, SFT, LoRA, DPO, PPO, GRPO, CISPO, uso de herramientas, Agentic RL y destilación.
+- Compatibilidad con Transformers, llama.cpp, vLLM y Ollama, además de una API compatible con OpenAI y una WebUI.
+- Conjuntos de datos, pesos y scripts de entrenamiento abiertos para aprender y reproducir resultados.
+
+**Inicio rápido**
+
+```bash
+git clone --depth 1 https://github.com/jingyaogong/minimind
+cd minimind
+pip install -r requirements.txt
+modelscope download --model gongjy/minimind-3 --local_dir ./minimind-3
+python eval_llm.py --load_from ./minimind-3
+```
+
+Para obtener instrucciones completas sobre entrenamiento, datos, evaluación y conversión de modelos, consulta la [documentación completa en inglés](./README_en.md).
+
+<a id="portuguese"></a>
+
+## Português
+
+MiniMind é um projeto de código aberto para treinar do zero um pequeno modelo de linguagem com cerca de 64 milhões de parâmetros. A etapa de ajuste fino supervisionado pode ser concluída em aproximadamente duas horas com uma única NVIDIA 3090, por um custo de locação próximo de 3 RMB.
+
+- Implementação PyTorch mínima e legível de arquiteturas Dense e MoE.
+- Pipeline completo: pré-treinamento, SFT, LoRA, DPO, PPO, GRPO, CISPO, uso de ferramentas, Agentic RL e destilação.
+- Compatibilidade com Transformers, llama.cpp, vLLM e Ollama, além de API compatível com OpenAI e WebUI.
+- Conjuntos de dados, pesos e scripts de treinamento abertos para aprendizado e reprodução dos resultados.
+
+**Início rápido**
+
+```bash
+git clone --depth 1 https://github.com/jingyaogong/minimind
+cd minimind
+pip install -r requirements.txt
+modelscope download --model gongjy/minimind-3 --local_dir ./minimind-3
+python eval_llm.py --load_from ./minimind-3
+```
+
+Para instruções completas sobre treinamento, dados, avaliação e conversão de modelos, consulte a [documentação completa em inglês](./README_en.md).
+
+<a id="german"></a>
+
+## Deutsch
+
+MiniMind ist ein Open-Source-Projekt, mit dem sich ein kleines Sprachmodell mit rund 64 Millionen Parametern von Grund auf trainieren lässt. Das überwachte Fine-Tuning kann auf einer einzelnen NVIDIA 3090 in ungefähr zwei Stunden abgeschlossen werden; die entsprechenden Mietkosten liegen bei etwa 3 RMB.
+
+- Kleine, gut lesbare PyTorch-Implementierung von Dense- und MoE-Architekturen.
+- Vollständige Pipeline: Pretraining, SFT, LoRA, DPO, PPO, GRPO, CISPO, Tool Use, Agentic RL und Distillation.
+- Kompatibel mit Transformers, llama.cpp, vLLM und Ollama; zusätzlich stehen eine OpenAI-kompatible API und eine WebUI bereit.
+- Offene Datensätze, Gewichte und Trainingsskripte zum Lernen und Reproduzieren der Ergebnisse.
+
+**Schnellstart**
+
+```bash
+git clone --depth 1 https://github.com/jingyaogong/minimind
+cd minimind
+pip install -r requirements.txt
+modelscope download --model gongjy/minimind-3 --local_dir ./minimind-3
+python eval_llm.py --load_from ./minimind-3
+```
+
+Vollständige Anleitungen zu Training, Datensätzen, Evaluation und Modellkonvertierung enthält die [englische Dokumentation](./README_en.md).
+
+<a id="persian"></a>
+
+<h2 dir="rtl">فارسی</h2>
+
+<p dir="rtl">MiniMind یک پروژهٔ متن‌باز برای آموزش یک مدل زبانی کوچک با حدود ۶۴ میلیون پارامتر از صفر است. مرحلهٔ تنظیم دقیق نظارت‌شده را می‌توان با یک کارت NVIDIA 3090 در حدود دو ساعت و با هزینهٔ اجاره‌ای نزدیک به ۳ یوان چین انجام داد.</p>
+
+<ul dir="rtl">
+  <li>پیاده‌سازی ساده و خوانای معماری‌های Dense و MoE با PyTorch.</li>
+  <li>زنجیرهٔ کامل آموزش شامل پیش‌آموزش، SFT، LoRA، DPO، PPO، GRPO، CISPO، استفاده از ابزار، Agentic RL و تقطیر مدل.</li>
+  <li>سازگار با Transformers، llama.cpp، vLLM و Ollama، همراه با API سازگار با OpenAI و WebUI.</li>
+  <li>داده‌ها، وزن‌ها و اسکریپت‌های آموزشی باز برای یادگیری و بازتولید نتایج.</li>
+</ul>
+
+<p dir="rtl"><strong>شروع سریع</strong></p>
+
+```bash
+git clone --depth 1 https://github.com/jingyaogong/minimind
+cd minimind
+pip install -r requirements.txt
+modelscope download --model gongjy/minimind-3 --local_dir ./minimind-3
+python eval_llm.py --load_from ./minimind-3
+```
+
+<p dir="rtl">برای دستورالعمل کامل آموزش، داده‌ها، ارزیابی و تبدیل مدل، <a href="./README_en.md">مستندات کامل انگلیسی</a> را ببینید.</p>
+
+<a id="japanese"></a>
+
+## 日本語
+
+MiniMind は、約 6,400 万パラメータの小規模言語モデルをゼロから学習できるオープンソースプロジェクトです。教師ありファインチューニングは、NVIDIA 3090 1枚で約2時間、レンタル費用約3人民元で完了できます。
+
+- Dense および MoE アーキテクチャを、読みやすい最小限の PyTorch コードで実装。
+- 事前学習、SFT、LoRA、DPO、PPO、GRPO、CISPO、Tool Use、Agentic RL、蒸留までを網羅。
+- Transformers、llama.cpp、vLLM、Ollama に対応し、OpenAI互換APIとWebUIも提供。
+- 学習と再現のためのデータセット、重み、学習スクリプトを公開。
+
+**クイックスタート**
+
+```bash
+git clone --depth 1 https://github.com/jingyaogong/minimind
+cd minimind
+pip install -r requirements.txt
+modelscope download --model gongjy/minimind-3 --local_dir ./minimind-3
+python eval_llm.py --load_from ./minimind-3
+```
+
+学習、データセット、評価、モデル変換の詳しい手順については、[英語版の完全なドキュメント](./README_en.md)を参照してください。
+
+<a id="korean"></a>
+
+## 한국어
+
+MiniMind는 약 6,400만 개의 파라미터를 가진 소형 언어 모델을 처음부터 학습할 수 있는 오픈 소스 프로젝트입니다. 지도 미세 조정 단계는 NVIDIA 3090 한 장에서 약 2시간, 약 3위안의 대여 비용으로 완료할 수 있습니다.
+
+- Dense 및 MoE 아키텍처를 간결하고 읽기 쉬운 PyTorch 코드로 구현합니다.
+- 사전 학습, SFT, LoRA, DPO, PPO, GRPO, CISPO, 도구 사용, Agentic RL, 모델 증류를 포함한 전체 파이프라인을 제공합니다.
+- Transformers, llama.cpp, vLLM, Ollama와 호환되며 OpenAI 호환 API와 WebUI를 제공합니다.
+- 학습과 결과 재현을 위한 데이터셋, 가중치, 학습 스크립트를 공개합니다.
+
+**빠른 시작**
+
+```bash
+git clone --depth 1 https://github.com/jingyaogong/minimind
+cd minimind
+pip install -r requirements.txt
+modelscope download --model gongjy/minimind-3 --local_dir ./minimind-3
+python eval_llm.py --load_from ./minimind-3
+```
+
+학습, 데이터셋, 평가, 모델 변환에 대한 전체 안내는 [영문 문서](./README_en.md)를 참고하세요.
+
+---
+
+These translations are community-maintained summaries. If you find an inaccurate phrase or want to maintain a complete translation, please open a pull request.

eval_llm.py

@@ -23,7 +23,7 @@ def init_model(args):
         model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
         if args.lora_weight != 'None':
             apply_lora(model)
-            load_lora(model, f'./{args.save_dir}/{args.lora_weight}_{args.hidden_size}.pth')
+            load_lora(model, f'./{args.save_dir}/{args.lora_weight}_{args.hidden_size}{moe_suffix}.pth')
     else:
         model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
     get_model_params(model, model.config)

scripts/serve_openai_api.py

@@ -40,7 +40,7 @@ def init_model(args):
         model.load_state_dict(torch.load(ckp, map_location=device), strict=True)
         if args.lora_weight != 'None':
             apply_lora(model)
-            load_lora(model, f'../{args.save_dir}/lora/{args.lora_weight}_{args.hidden_size}.pth')
+            load_lora(model, f'../{args.save_dir}/{args.lora_weight}_{args.hidden_size}{moe_suffix}.pth')
     else:
         model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
     print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M(illion)')

trainer/train_agent.py

@@ -240,12 +240,10 @@ def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, rewa
 
 # ================================ 工具与 Reward = End ================================
 def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model=None, start_step=0, wandb=None, use_sglang=False):
-    last_step = start_step
     for step, batch in enumerate(loader, start=start_step + 1):
         messages_batch = batch['messages']
         tools_batch = batch['tools']
         gt_batch = batch['gt']
-        last_step = step
 
         with torch.no_grad():
             completions, contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch, turn_outputs_batch, unfinished_batch = rollout_batch(rollout_engine, tokenizer, messages_batch, tools_batch, args.num_generations, max_turns=3, max_new_tokens=args.max_gen_len, thinking_ratio=args.thinking_ratio, device=args.device)
@@ -331,7 +329,7 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
         loss = (policy_loss + aux_loss) / args.accumulation_steps
         loss.backward()
 
-        if step % args.accumulation_steps == 0:
+        if step % args.accumulation_steps == 0 or step == iters:
             if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
             optimizer.step(); scheduler.step(); optimizer.zero_grad()
 
@@ -365,10 +363,6 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
         del per_token_logps, ref_per_token_logps
         del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
 
-    if last_step > start_step and last_step % args.accumulation_steps != 0:
-        if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-        optimizer.step(); scheduler.step(); optimizer.zero_grad()
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Agent RL")

trainer/train_distillation.py

@@ -38,14 +38,12 @@ def distillation_loss(student_logits, teacher_logits, temperature=1.0, reduction
 
 def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_step=0, wandb=None, alpha=0.0, temperature=1.0):
     start_time = time.time()
-    last_step = start_step
     
     if teacher_model is not None:
         teacher_model.eval()
         teacher_model.requires_grad_(False)
 
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
-        last_step = step
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
         loss_mask = (labels[..., 1:] != -100).float()
@@ -94,7 +92,7 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st
 
         scaler.scale(loss).backward()
 
-        if step % args.accumulation_steps == 0:
+        if step % args.accumulation_steps == 0 or step == iters:
             scaler.unscale_(optimizer)
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
             scaler.step(optimizer)
@@ -135,13 +133,6 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st
 
         del input_ids, labels, loss_mask, res, student_logits, ce_loss, distill_loss, loss
 
-    if last_step > start_step and last_step % args.accumulation_steps != 0:
-        scaler.unscale_(optimizer)
-        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad(set_to_none=True)
-
 
 if __name__ == "__main__":
     # 模拟用moe模型蒸馏dense模型，也可以用更大teacher_hidden_size模型蒸馏更小student_hidden_size的

trainer/train_dpo.py

@@ -52,10 +52,8 @@ def dpo_loss(ref_log_probs, policy_log_probs, mask, beta):
 
 def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=None, beta=0.1):
     start_time = time.time()
-    last_step = start_step
 
     for step, batch in enumerate(loader, start=start_step + 1):
-        last_step = step
         x_chosen = batch['x_chosen'].to(args.device)
         x_rejected = batch['x_rejected'].to(args.device)
         y_chosen = batch['y_chosen'].to(args.device)
@@ -86,7 +84,7 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=
 
         scaler.scale(loss).backward()
 
-        if step % args.accumulation_steps == 0:
+        if step % args.accumulation_steps == 0 or step == iters:
             scaler.unscale_(optimizer)
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
             scaler.step(optimizer)
@@ -120,13 +118,6 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=
         del x_chosen, x_rejected, y_chosen, y_rejected, mask_chosen, mask_rejected, x, y, mask
         del ref_outputs, ref_logits, ref_log_probs, outputs, logits, policy_log_probs, loss
 
-    if last_step > start_step and last_step % args.accumulation_steps != 0:
-        scaler.unscale_(optimizer)
-        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad(set_to_none=True)
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind DPO (Direct Preference Optimization)")

trainer/train_full_sft.py

@@ -23,11 +23,9 @@ warnings.filterwarnings('ignore')
 
 def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
     start_time = time.time()
-    last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
-        last_step = step
         lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -39,7 +37,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
 
         scaler.scale(loss).backward()
 
-        if step % args.accumulation_steps == 0:
+        if step % args.accumulation_steps == 0 or step == iters:
             scaler.unscale_(optimizer)
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
 
@@ -73,13 +71,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
 
         del input_ids, labels, res, loss
 
-    if last_step > start_step and last_step % args.accumulation_steps != 0:
-        scaler.unscale_(optimizer)
-        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad(set_to_none=True)
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Full SFT")

trainer/train_grpo.py

@@ -144,7 +144,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
         loss = (policy_loss + aux_loss) / args.accumulation_steps  # scalar
         loss.backward()
 
-        if step % args.accumulation_steps == 0:
+        if step % args.accumulation_steps == 0 or step == iters:
             if args.grad_clip > 0:
                 torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
             optimizer.step()
@@ -195,13 +195,6 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
         del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
         del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask, completion_pad_mask, prompt_lens, logp_pos
 
-    if step > start_step and step % args.accumulation_steps != 0:
-        if args.grad_clip > 0:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-        optimizer.step()
-        scheduler.step()
-        optimizer.zero_grad()
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind GRPO (Group Relative Policy Optimization)")

trainer/train_lora.py

@@ -24,11 +24,9 @@ warnings.filterwarnings('ignore')
 
 def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
     start_time = time.time()
-    last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
-        last_step = step
         lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -40,7 +38,7 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
 
         scaler.scale(loss).backward()
 
-        if step % args.accumulation_steps == 0:
+        if step % args.accumulation_steps == 0 or step == iters:
             scaler.unscale_(optimizer)
             torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
             scaler.step(optimizer)
@@ -68,13 +66,6 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
 
         del input_ids, labels, res, loss
 
-    if last_step > start_step and last_step % args.accumulation_steps != 0:
-        scaler.unscale_(optimizer)
-        torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad(set_to_none=True)
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind LoRA Fine-tuning")
     parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")

trainer/train_pretrain.py

@@ -23,11 +23,9 @@ warnings.filterwarnings('ignore')
 
 def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
     start_time = time.time()
-    last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
         input_ids = input_ids.to(args.device)
         labels = labels.to(args.device)
-        last_step = step
         lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -39,7 +37,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
 
         scaler.scale(loss).backward()
 
-        if step % args.accumulation_steps == 0:
+        if step % args.accumulation_steps == 0 or step == iters:
             scaler.unscale_(optimizer)
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
 
@@ -72,13 +70,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
 
         del input_ids, labels, res, loss
 
-    if last_step > start_step and last_step % args.accumulation_steps != 0:
-        scaler.unscale_(optimizer)
-        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-        scaler.step(optimizer)
-        scaler.update()
-        optimizer.zero_grad(set_to_none=True)
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Pretraining")