update lr

This commit is contained in:
gongjy 2025-02-11 23:53:48 +08:00
parent d2f5ef4355
commit bab480073e

View File

@ -122,7 +122,7 @@ def init_distributed_mode():
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser(description="MiniMind Full SFT") parser = argparse.ArgumentParser(description="MiniMind Full SFT")
parser.add_argument("--out_dir", type=str, default="out") parser.add_argument("--out_dir", type=str, default="out")
parser.add_argument("--epochs", type=int, default=6) parser.add_argument("--epochs", type=int, default=1)
parser.add_argument("--batch_size", type=int, default=32) parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--learning_rate", type=float, default=5e-5) parser.add_argument("--learning_rate", type=float, default=5e-5)
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")