Upload src/bluscriptCP.cpp with huggingface_hub

src/bluscriptCP.cpp

@@ -83,6 +83,7 @@
 #include "DataParallel.hpp"                     // DataParallel, DDP_Options
 #include "dist/zero/include/ShardedOptimizer.h"           // owt::zero::ShardedOptimizer (ZeRO-1)
 #include "dist/zero/include/SymmetricShardedOptimizer.h"  // ZeRO-1 v2 (reduce-scatter/all-gather)
+#include "dist/zero/include/ShardedMuon.h"              // owt::zero::ShardedMuon (ZeRO-1 over whole tensors)
 #include "dist/zero/include/ZeroOptimizerBase.h"           // common surface: v1/v2 interchangeable
 
 // Context parallelism
@@ -148,8 +149,9 @@ struct ModelConfig {
     float   grad_clip    = 1.0f;
     float   weight_decay = 0.1f;
 
-    float   max_lr       = 6e-4f;
-    float   min_lr       = 6e-5f;                        // 10% of max_lr
+    float   max_lr       = 6e-4f;                        // CP_MAX_LR overrides
+    float   min_lr       = 6e-5f;                        // CP_MIN_LR overrides (10% of max_lr)
+    bool    use_muon     = false;                        // CP_OPTIMIZER=muon
     int     warmup_steps = 91;
     // CP_BF16=1 makes the Linears (wQ/wK/wV/wO, gate_up/down, lm_head matmul)
     // bf16 with fp32 norms, fp32 residual stream and fp32 logits -- the scheme
@@ -982,6 +984,20 @@ int main(int argc, char** argv) {
     cfg.max_steps    = static_cast<int>(env_i64("CP_MAX_STEPS", cfg.max_steps));
     cfg.warmup_steps = static_cast<int>(env_i64("CP_WARMUP", cfg.warmup_steps));
     cfg.stable_steps = static_cast<int>(env_i64("CP_STABLE", cfg.stable_steps));
+    // LR BOUNDS FROM ENV (added 2026-09-18). max_lr/min_lr were hardcoded 6e-4 /
+    // 6e-5 and NOTHING read CP_MAX_LR or CP_MIN_LR -- every script in this
+    // project that set them was setting a variable the trainer ignored, and the
+    // floor could not be moved at all. Same failure family as CP_OPTIMIZER below.
+    if (const char* e = std::getenv("CP_MAX_LR")) cfg.max_lr = std::atof(e);
+    if (const char* e = std::getenv("CP_MIN_LR")) cfg.min_lr = std::atof(e);
+    // CP_OPTIMIZER=muon selects Muon for the >=2-D weight matrices with an aux
+    // AdamW path for everything else, mirroring bluscript_zero_ckpt.cpp's
+    // BLU_OPTIMIZER=muon. Until now this flag was READ NOWHERE: every CP run
+    // that passed CP_OPTIMIZER=muon in fact trained with AdamW.
+    cfg.use_muon = [] {
+        const char* v = std::getenv("CP_OPTIMIZER");
+        return v && (std::string(v) == "muon" || std::string(v) == "MUON");
+    }();
     cfg.bf16         = env_i64("CP_BF16", cfg.bf16 ? 1 : 0) != 0;
     cfg.bf16_mask    = static_cast<int>(env_i64("CP_BF16_MASK", cfg.bf16_mask));
     // Defaults ON when bf16 is requested at world_size>1, because DDP is broken
@@ -1716,8 +1732,34 @@ int main(int argc, char** argv) {
             // direction. Reduce in fp32 when that is suspected.
             zcfg.reduce_dtype = owt::zero::ReduceDtype::FP32;
         }
+        // Muon config mirrors bluscript_zero_ckpt.cpp:1319-1347 exactly, with ONE
+        // deliberate difference: lr and adamw_lr are both set to cfg.max_lr so the
+        // ratio set_lr() preserves is 1:1 and BOTH paths run at the same learning
+        // rate. Pretraining used 8e-4 / 6e-4 (a 4:3 Muon:AdamW ratio); here the
+        // ablation calls for one shared value.
+        owt::zero::MuonConfig mcfg;
+        mcfg.collective         = owt::zero::MuonCollective::REDUCE_SCATTER_ALL_GATHER;
+        mcfg.target_buckets     = static_cast<int>(env_i64("CP_MUON_BUCKETS", 6));
+        mcfg.lr                 = cfg.max_lr;
+        mcfg.momentum           = 0.95f;
+        mcfg.weight_decay       = cfg.weight_decay;
+        mcfg.ns_steps           = 5;
+        mcfg.nesterov           = true;
+        mcfg.adamw_lr           = cfg.max_lr;      // 1:1 with the Muon path
+        mcfg.adamw_beta1        = zcfg.beta1;
+        mcfg.adamw_beta2        = zcfg.beta2;
+        mcfg.adamw_eps          = zcfg.eps;
+        mcfg.adamw_weight_decay = cfg.weight_decay;
+        mcfg.exclude_1d_from_weight_decay = true;
+        mcfg.max_grad_norm      = zcfg.max_grad_norm;
+        mcfg.reduce_dtype       = zcfg.reduce_dtype;
+        mcfg.init_broadcast     = zcfg.init_broadcast;
+        mcfg.verbose            = is_master;   // otherwise Muon runs SILENTLY and no
+                                               // log line distinguishes it from AdamW
         for (auto& kv : params_by_dtype) {
-            if (cfg.zero_v2)
+            if (cfg.use_muon)
+                zopts.push_back(std::make_unique<owt::zero::ShardedMuon>(kv.second, pg, mcfg));
+            else if (cfg.zero_v2)
                 zopts.push_back(std::make_unique<owt::zero::SymmetricShardedOptimizer>(kv.second, pg, zcfg));
             else
                 zopts.push_back(std::make_unique<owt::zero::ShardedOptimizer>(kv.second, pg, zcfg));
@@ -1725,8 +1767,14 @@ int main(int argc, char** argv) {
             zgroup_label.push_back(OwnTensor::get_dtype_name(kv.first) + "/" + std::to_string(n));
         }
         if (is_master) {
-            std::cout << "[ZeRO-1] variant=" << (cfg.zero_v2 ? "v2 reduce-scatter/all-gather"
-                                                            : "v1 reduce/broadcast")
+            // Report what was actually BUILT. This line used to print cfg.zero_v2
+            // unconditionally, so a Muon run was indistinguishable from an AdamW
+            // one in the log -- the exact silent-failure shape this file keeps
+            // hitting (CP_OPTIMIZER, CP_MAX_LR, CP_MIN_LR were all no-ops).
+            std::cout << "[ZeRO-1] optimizer=" << (cfg.use_muon ? "MUON+adamw-aux" : "adamw")
+                      << "  variant=" << (cfg.use_muon ? "muon whole-tensor shard"
+                                          : (cfg.zero_v2 ? "v2 reduce-scatter/all-gather"
+                                                         : "v1 reduce/broadcast"))
                       << "  groups=" << zopts.size();
             for (auto& g : zgroup_label) std::cout << "  [" << g << "]";
             std::cout << "\n[ZeRO-1] DataParallel is DISABLED (both would double-reduce)"