[Misc.] Roll back the initial std range. Use larger learning rate by default

Author: sustcsonglin

README.md

@@ -58,12 +58,12 @@ Here's an example of training a 340M FLA Transformer model with a LLaMA-like arc
 ```sh
 bash train.sh \
   --job.config_file flame/models/fla.toml \
-  --job.dump_folder exp/transformer-340M-4K-10B/batch1.seqlen65536.context4096.warmup1024.update1.steps20480.lr3e-4.cosine \
+  --job.dump_folder exp/transformer-340M-4K-10B/batch1.seqlen65536.context4096.warmup1024.update1.steps20480.lr1e-3.cosine \
   --model.config configs/transformer_340M.json \
   --model.tokenizer_path fla-hub/transformer-1.3B-100B \
   --optimizer.name AdamW \
   --optimizer.eps 1e-15 \
-  --optimizer.lr 3e-4 \
+  --optimizer.lr 1e-3 \
   --lr_scheduler.warmup_steps 1024 \
   --lr_scheduler.lr_min 0.1 \
   --lr_scheduler.decay_type cosine \
@@ -92,7 +92,7 @@ You can specify the number of GPUs by setting the environment variable `NGPU`, w
 **For single-GPU debugging, set `NGPU=1`.**
 
 We provide several [config files](https://github.com/fla-org/flame/tree/main/configs) for different models.
-By default, the learning rate is set to 3e-4 with a cosine scheduler. Other schedulers, such as WSD (wsd), are also supported.
+By default, the learning rate is set to 1e-3 with a cosine scheduler. Other schedulers, such as WSD (wsd), are also supported.
 
 **Key parameters:**
 - `--lr_scheduler.decay_ratio`: The proportion of the steps allocated to the decay phase. The learning rate will remain stable after the warmup period and only start decaying during the last `decay_ratio` portion of the total training steps, which is known as the Warmup-Stable-Decay (WSD) schedule.

configs/delta_net_1B.json

@@ -11,7 +11,7 @@
     "hidden_act": "swish",
     "hidden_ratio": 4,
     "hidden_size": 2048,
-    "initializer_range": 0.006,
+    "initializer_range": 0.02,
     "intermediate_size": null,
     "model_type": "delta_net",
     "norm_eps": 1e-06,

configs/delta_net_340M.json

@@ -9,7 +9,7 @@
     "hidden_act": "swish",
     "hidden_ratio": 4,
     "hidden_size": 1024,
-    "initializer_range": 0.006,
+    "initializer_range": 0.02,
     "intermediate_size": null,
     "model_type": "delta_net",
     "norm_eps": 1e-06,

configs/gla_340M.json

@@ -10,7 +10,7 @@
   "hidden_act": "swish",
   "hidden_ratio": 4,
   "hidden_size": 1024,
-  "initializer_range": 0.006,
+  "initializer_range": 0.02,
   "intermediate_size": null,
   "model_type": "gla",
   "num_heads": 4,

configs/gla_7B.json

@@ -10,7 +10,7 @@
     "hidden_act": "swish",
     "hidden_ratio": 4,
     "hidden_size": 4096,
-    "initializer_range": 0.006,
+    "initializer_range": 0.02,
     "intermediate_size": 11008,
     "model_type": "gla",
     "norm_eps": 1e-06,

configs/gsa_340M.json

@@ -12,7 +12,7 @@
     "hidden_act": "swish",
     "hidden_ratio": 4,
     "hidden_size": 1024,
-    "initializer_range": 0.006,
+    "initializer_range": 0.02,
     "intermediate_size": null,
     "model_type": "gsa",
     "num_heads": 4,

configs/hgrn2_340M.json

@@ -8,7 +8,7 @@
     "hidden_act": "swish",
     "hidden_ratio": 4,
     "hidden_size": 1024,
-    "initializer_range": 0.006,
+    "initializer_range": 0.02,
     "intermediate_size": null,
     "model_type": "hgrn2",
     "num_heads": 8,

configs/transformer_1B.json

@@ -8,7 +8,7 @@
     "hidden_act": "swish",
     "hidden_ratio": 4,
     "hidden_size": 2048,
-    "initializer_range": 0.006,
+    "initializer_range": 0.02,
     "intermediate_size": null,
     "max_position_embeddings": 8192,
     "model_type": "transformer",

configs/transformer_340M.json

@@ -6,7 +6,7 @@
     "fuse_norm": true,
     "hidden_act": "swish",
     "hidden_size": 1024,
-    "initializer_range": 0.006,
+    "initializer_range": 0.02,
     "max_position_embeddings": 8192,
     "model_type": "transformer",
     "num_heads": 16,

configs/transformer_7B.json

@@ -7,7 +7,7 @@
     "hidden_act": "swish",
     "hidden_ratio": 4,
     "hidden_size": 4096,
-    "initializer_range": 0.006,
+    "initializer_range": 0.02,
     "intermediate_size": 14336,
     "model_type": "transformer",
     "norm_eps": 1e-06,