Add 350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415

.gitattributes

@@ -58,3 +58,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 60m/soap_lr3e_3_b1_0_9_b2_0_95_eps_1e_8_A100_ppl_29_4706_20260416_193855/wandb/offline-run-20260416_193930-flol3ksy/run-flol3ksy.wandb filter=lfs diff=lfs merge=lfs -text
 60m/sophia_lr2e_4_b1_0_9_b2_0_99_eps_1e_8_A100_ppl_36_2695_20260416_193855/wandb/offline-run-20260416_193930-x2s18q4b/run-x2s18q4b.wandb filter=lfs diff=lfs merge=lfs -text
 130m/adabelief_lr1e_3_b1_0_9_b2_0_999_eps_1e_16_A100_ppl_23_4537_20260417_181954/wandb/offline-run-20260417_221600-a8vnu42p/run-a8vnu42p.wandb filter=lfs diff=lfs merge=lfs -text
+350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/run-mawza3ul.wandb filter=lfs diff=lfs merge=lfs -text

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/model_60000/config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dtype": "bfloat16",
+  "eos_token_id": 1,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 2736,
+  "max_position_embeddings": 2048,
+  "max_sequence_length": 1024,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 16,
+  "pad_token_id": -1,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.3",
+  "use_cache": true,
+  "vocab_size": 32000
+}

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/model_60000/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae4737309e6d86fd1a6704986f79d99c5701b932b40a55b1defe52265b26314c
+size 735967792

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/model_60000/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fae0081c25762de738fa1ee41275b545fbcb33e26dbbc712ac740dbffe6185bf
+size 1824395851

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/model_60000/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc0e795939cf0c432c6675fb69805b3f4beaadf0f7e4d45715ecb5cc46766541
+size 736040495

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/model_60000/training_state.json

@@ -0,0 +1,7 @@
+{
+    "global_step": 60000,
+    "update_step": 60000,
+    "tokens_seen": 5997831592,
+    "tokens_seen_before": 5997732940,
+    "update_time": 1.0997016429901123
+}

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb.json

@@ -0,0 +1,3 @@
+{
+    "wandb_id": "mawza3ul"
+}

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb/debug-internal.log

@@ -0,0 +1,15 @@
+{"time":"2026-04-23T20:24:22.4996429+08:00","level":"INFO","msg":"wandb-core"}
+{"time":"2026-04-23T20:24:22.499901972+08:00","level":"INFO","msg":"stream: starting","core version":"0.26.0"}
+{"time":"2026-04-23T20:24:22.63528813+08:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
+{"time":"2026-04-23T20:24:22.635312622+08:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
+{"time":"2026-04-23T20:24:22.635332006+08:00","level":"INFO","msg":"stream: created new stream","id":"mawza3ul"}
+{"time":"2026-04-23T20:24:22.635539387+08:00","level":"INFO","msg":"handler: started"}
+{"time":"2026-04-23T20:24:22.636204984+08:00","level":"INFO","msg":"stream: started"}
+{"time":"2026-04-23T20:24:22.636292599+08:00","level":"INFO","msg":"writer: started","stream_id":"mawza3ul"}
+{"time":"2026-04-23T20:24:22.636304487+08:00","level":"INFO","msg":"sender: started"}
+{"time":"2026-04-23T20:24:22.637327692+08:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
+{"time":"2026-04-23T20:24:22.637343161+08:00","level":"WARN","msg":"runupserter: server does not expand metric globs but the x_server_side_expand_glob_metrics setting is set; ignoring"}
+{"time":"2026-04-24T01:14:24.730723575+08:00","level":"INFO","msg":"stream: closing"}
+{"time":"2026-04-24T01:14:24.753028533+08:00","level":"INFO","msg":"handler: closed"}
+{"time":"2026-04-24T01:14:24.753387495+08:00","level":"INFO","msg":"sender: closed"}
+{"time":"2026-04-24T01:14:24.753400977+08:00","level":"INFO","msg":"stream: closed"}

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb/debug.log

@@ -0,0 +1,22 @@
+2026-04-23 20:24:22,097 INFO    MainThread:342 [wandb_setup.py:_flush():81] Current SDK version is 0.26.0
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_setup.py:_flush():81] Configure stats pid to 342
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:setup_run_log_directory():721] Logging user logs to exp_remain_h200/work_dirs/350m/train_350m_conda_lr1e_2_scale0_25_rank256_gap2000_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/logs/debug.log
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:setup_run_log_directory():722] Logging internal logs to exp_remain_h200/work_dirs/350m/train_350m_conda_lr1e_2_scale0_25_rank256_gap2000_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/logs/debug-internal.log
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:init():848] calling init triggers
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:init():853] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:init():896] starting backend
+2026-04-23 20:24:22,494 INFO    MainThread:342 [wandb_init.py:init():911] sending inform_init request
+2026-04-23 20:24:22,498 INFO    MainThread:342 [wandb_init.py:init():919] backend started and connected
+2026-04-23 20:24:22,498 INFO    MainThread:342 [wandb_init.py:init():989] updated telemetry
+2026-04-23 20:24:22,526 INFO    MainThread:342 [wandb_init.py:init():1013] communicating run to backend with 90.0 second timeout
+2026-04-23 20:24:22,638 INFO    MainThread:342 [wandb_init.py:init():1058] starting run threads in backend
+2026-04-23 20:24:22,712 INFO    MainThread:342 [wandb_run.py:_console_start():2542] atexit reg
+2026-04-23 20:24:22,712 INFO    MainThread:342 [wandb_run.py:_redirect():2391] redirect: wrap_raw
+2026-04-23 20:24:22,712 INFO    MainThread:342 [wandb_run.py:_redirect():2460] Wrapping output streams.
+2026-04-23 20:24:22,712 INFO    MainThread:342 [wandb_run.py:_redirect():2483] Redirects installed.
+2026-04-23 20:24:22,714 INFO    MainThread:342 [wandb_init.py:init():1098] run started, returning control to user process
+2026-04-23 20:24:49,867 INFO    MainThread:342 [wandb_run.py:_config_callback():1403] config_cb None None {'model_config': 'configs/llama_350m.json', 'exp_config': 'exp_v2/configs/llama_350m.json', 'eval_every': 1000, 'save_every': 60000, 'dtype': 'bfloat16', 'seed': 0, 'compile': True, 'dynamo_suppress_errors': True, 'dynamo_cache_limit': 10000, 'memory_cleanup_frequency': 10000, 'resume_step': None, 'restore_optimizer': False, 'continue_from': None, 'single_gpu': False, 'save_dir': 'exp_remain_h200/work_dirs/350m/train_350m_conda_lr1e_2_scale0_25_rank256_gap2000_20260423_202415', 'use_hf_model': False, 'workers': 12, 'batch_size': 128, 'gradient_accumulation': 1, 'total_batch_size': 512, 'warmup_steps': 6000, 'num_training_steps': 60000, 'max_train_tokens': None, 'optimizer': 'conda', 'max_length': 256, 'scheduler': 'cosine', 'min_lr_ratio': 0.1, 'weight_decay': 0.0, 'grad_clipping': 0.0, 'activation_checkpointing': False, 'data_path': '/mnt/shared-storage-gpfs2/finebio-shared/optimizer/dataset/C4/en', 'data_name': 'en', 'tags': None, 'name': 'test', 'project': 'test', 'unset_wandb': False, 'entity': None, 'wandb_dir': 'exp_remain_h200/work_dirs/350m/train_350m_conda_lr1e_2_scale0_25_rank256_gap2000_20260423_202415', 'beta1': 0.9, 'beta2': 0.99, 'beta3': 0.99, 'eps': 1e-08, 'rank': 256, 'update_proj_gap': 2000, 'galore_scale': 1.0, 'proj_type': 'std', 'proj_quant': False, 'proj_bits': 8, 'proj_group_size': 256, 'weight_quant': False, 'weight_bits': 8, 'weight_group_size': 256, 'stochastic_round': False, 'simulation': False, 'cos_threshold': 1, 'gamma_proj': 2, 'queue_size': 5, 'proj': 'random', 'scale_type': 'channel', 'apollo_scale': 0.25, 'scale_front': False, 'n_clusters': 3, 'scale_update_freq': 500, 'scale_level': '1,0,1,1', 'scale_bound': None, 'metric': 'mean', 'align_grad': False, 'dim': 4096, 'n_heads': 32, 'muon_ns_steps': 5, 'muon_momentum': 0.95, 'nproc_per_node': 4, 'max_lr': 0.01, 'total_params_M': 367.96928, 'dataset': 'c4', 'model': {'vocab_size': 32000, 'max_position_embeddings': 2048, 'hidden_size': 1024, 'intermediate_size': 2736, 'num_hidden_layers': 24, 'num_attention_heads': 16, 'num_key_value_heads': 16, 'hidden_act': 'silu', 'initializer_range': 0.02, 'rms_norm_eps': 1e-06, 'pretraining_tp': 1, 'use_cache': True, 'rope_theta': 10000.0, 'rope_scaling': None, 'attention_bias': False, 'attention_dropout': 0.0, 'mlp_bias': False, 'head_dim': 64, 'return_dict': True, 'output_hidden_states': False, 'torchscript': False, 'dtype': None, 'pruned_heads': {}, 'tie_word_embeddings': False, 'chunk_size_feed_forward': 0, 'is_encoder_decoder': False, 'is_decoder': False, 'cross_attention_hidden_size': None, 'add_cross_attention': False, 'tie_encoder_decoder': False, 'architectures': ['LLaMAForCausalLM'], 'finetuning_task': None, 'id2label': {0: 'LABEL_0', 1: 'LABEL_1'}, 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, 'task_specific_params': None, 'problem_type': None, 'tokenizer_class': None, 'prefix': None, 'bos_token_id': 0, 'pad_token_id': -1, 'eos_token_id': 1, 'sep_token_id': None, 'decoder_start_token_id': None, 'max_length': 20, 'min_length': 0, 'do_sample': False, 'early_stopping': False, 'num_beams': 1, 'temperature': 1.0, 'top_k': 50, 'top_p': 1.0, 'typical_p': 1.0, 'repetition_penalty': 1.0, 'length_penalty': 1.0, 'no_repeat_ngram_size': 0, 'encoder_no_repeat_ngram_size': 0, 'bad_words_ids': None, 'num_return_sequences': 1, 'output_scores': False, 'return_dict_in_generate': False, 'forced_bos_token_id': None, 'forced_eos_token_id': None, 'remove_invalid_values': False, 'exponential_decay_length_penalty': None, 'suppress_tokens': None, 'begin_suppress_tokens': None, 'num_beam_groups': 1, 'diversity_penalty': 0.0, '_name_or_path': 'configs/llama_350m.json', 'transformers_version': '4.57.3', 'max_sequence_length': 1024, 'model_type': 'llama', 'tf_legacy_loss': False, 'use_bfloat16': False, 'output_attentions': False}, 'world_size': 4, 'device': 'cuda:0'}
+2026-04-24 01:14:24,727 INFO    wandb-AsyncioManager-main:342 [service_client.py:_forward_responses():134] Reached EOF.
+2026-04-24 01:14:24,730 INFO    wandb-AsyncioManager-main:342 [mailbox.py:close():155] Closing mailbox, abandoning 0 handles.

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/files/SAC/torchrun_main.py

@@ -0,0 +1,612 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import os
+import time
+import json
+
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+
+from tqdm import tqdm
+from loguru import logger
+
+import transformers
+
+transformers.logging.set_verbosity_error()
+
+import wandb
+
+from utils.argparse import parse_args
+from utils.setup import getting_svd_cnt, set_seed, setup_model, saving_model_weight, load_model_weight
+from utils.optimizer_factory import setup_optimization
+from utils.eval import evaluate_model
+from utils.dataloader import setup_dataset
+from utils.modeling_llama import LlamaForCausalLM
+from utils.fake_quantization import QLinear
+from utils.quantization import QScaleLinear
+from opt.hybrid import HybridOptimizer, HybridScheduler
+
+
+def main(args):
+    import torch
+    ############ Setup random seed ############
+    set_seed(args)
+
+    ############ Setup DDP environment ############
+    assert "LOCAL_RANK" in os.environ, "torchrun should set LOCAL_RANK"
+    global_rank = int(os.environ["RANK"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    torch.cuda.set_device(local_rank)
+
+    logger.info(f"Global rank {global_rank}, local rank {local_rank}, device: {torch.cuda.current_device()}")
+    dist.init_process_group(backend="nccl", rank=global_rank, world_size=world_size)
+
+    logger.info("Process group initialized")
+    device = f"cuda:{local_rank}"
+
+    if global_rank != 0:
+        logger.remove()  # turn off logger
+
+    logger.info(f"Using dist with rank {global_rank} (only rank 0 will log)")
+    logger.info("*" * 40)
+    logger.info(f"Starting training with the arguments")
+    for k, v in vars(args).items():
+        logger.info(f"{k:30} {v}")
+    logger.info("*" * 40)
+
+    ############ Initialize wandb without config (it is passed later) ############
+    if (not args.unset_wandb) and global_rank == 0:
+        if args.entity is None:
+            os.environ['WANDB_MODE'] = 'offline'
+        # Set wandb directory for offline mode
+        wandb_dir = getattr(args, 'wandb_dir', None) if getattr(args, 'wandb_dir', None) is not None else args.save_dir
+        if getattr(args, 'wandb_dir', None) is not None:
+            logger.info(f"Wandb directory set to: {wandb_dir}")
+        wandb.init(project=args.project, name=args.name, entity=args.entity, dir=wandb_dir)
+
+    ############ Setup training data ############
+    if args.total_batch_size is not None:
+        if args.gradient_accumulation is None:
+            assert args.total_batch_size % world_size == 0, "total_batch_size must be divisible by world_size"
+            args.gradient_accumulation = args.total_batch_size // (args.batch_size * world_size)
+            assert args.gradient_accumulation > 0, "gradient_accumulation must be greater than 0"
+
+    assert (
+        args.gradient_accumulation * args.batch_size * world_size == args.total_batch_size
+    ), "gradient_accumulation * batch_size * world_size must be equal to total_batch_size"
+
+    dataloader, tokenizer = setup_dataset(args, global_rank, world_size)
+
+    ############ Initialize model ############
+    model_config, model = setup_model(args)
+    # Ensure model has generation_config (fix for transformers version compatibility)
+    if not hasattr(model, 'generation_config') or model.generation_config is None:
+        from transformers import GenerationConfig
+        model.generation_config = GenerationConfig()
+    model.generation_config.pad_token_id = tokenizer.pad_token_id
+
+    ############ Resuming from checkpoints ############
+    global_step = 0
+    update_step = 0
+    beginning_step = 0
+    tokens_seen = 0
+    tokens_seen_before = 0
+
+    # identifying checkpointing
+    if args.continue_from is not None and os.path.exists(args.continue_from):
+        # searching the latest checkpoints
+        checkpoint_path_list = os.listdir(args.continue_from)
+        checkpoint_path_list = [int(x.split("_")[-1]) for x in checkpoint_path_list if x.startswith("model_")]
+        if len(checkpoint_path_list) > 0:
+            logger.info("Find Checkpoints", checkpoint_path_list)
+            beginning_step = max(checkpoint_path_list)
+            if args.resume_step is not None:
+                beginning_step = args.resume_step
+            args.continue_from = os.path.join(args.continue_from, f"model_{beginning_step}")
+            logger.info("Continue from", args.continue_from)
+        else:
+            logger.warning(f"Did not find any checkpoints in {args.continue_from}")
+            args.continue_from = None
+
+    # resuming from checkpointing
+    if args.continue_from is not None:
+        logger.info("*" * 40)
+        logger.info(f"Loading model from {args.continue_from}")
+        checkpoint_path = os.path.join(args.continue_from, "pytorch_model.bin")
+        if os.path.exists(checkpoint_path):
+            load_model_weight(model, checkpoint_path, args)
+            logger.info(f"Model successfully loaded (strict=False policy)")
+        else:
+            # Try safetensors format
+            checkpoint_path = os.path.join(args.continue_from, "model.safetensors")
+            if os.path.exists(checkpoint_path):
+                from safetensors import safe_open
+                tensors = {}
+                with safe_open(checkpoint_path, framework="pt", device=0) as f:
+                    for k in f.keys():
+                        tensors[k] = f.get_tensor(k)
+                        print(k, tensors[k].shape)
+                ret = model.load_state_dict(tensors, strict=False)
+                logger.info(f"Model successfully loaded from safetensors (strict=False policy)", ret)
+            else:
+                logger.warning(f"No model checkpoint found in {args.continue_from}")
+
+        if os.path.exists(os.path.join(args.continue_from, "training_state.json")):
+            logger.info(
+                f"Loading training state like global_step, update_step, and tokens_seen from {args.continue_from}"
+            )
+            with open(os.path.join(args.continue_from, "training_state.json")) as f:
+                _old_state = json.load(f)
+            global_step = _old_state["global_step"]
+            update_step = _old_state["update_step"]
+            tokens_seen = _old_state["tokens_seen"]
+            tokens_seen_before = _old_state["tokens_seen_before"]
+            logger.info(f"global_step       : {global_step}")
+            logger.info(f"update_step       : {update_step}")
+            logger.info(f"tokens_seen       : {tokens_seen}")
+            logger.info(f"tokens_seen_before: {tokens_seen_before}")
+            logger.info(f"Will train for {args.num_training_steps - update_step} update steps")
+        else:
+            logger.warning(f"Did not find training state in {args.continue_from}, global step will start from zero")
+        logger.info("*" * 40)
+
+    ############ Setup model ############
+    if args.dtype in ["bf16", "bfloat16"]:
+        model = model.to(dtype=torch.bfloat16)
+    model = model.to(device=device)
+
+    for _, module in model.named_modules():
+        if isinstance(module, QScaleLinear):
+            weight_device = module.weight.device
+            module.weight.scales = module.weight.scales.to(device=weight_device)
+            module.weight.zeros = module.weight.zeros.to(device=weight_device)
+
+    n_total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = [p for p in model.parameters() if p.requires_grad]
+    trainable_params_int8 = [p for p in model.parameters() if hasattr(p, "group_size")]
+
+    ############ Initialize wandb ############
+    run_config = dict(vars(args))
+    run_config.update(
+        {
+            "max_lr": run_config.pop("lr"),  # rename lr to max_lr to avoid conflicts with scheduler
+            "total_params_M": n_total_params / 1_000_000,
+            "dataset": "c4",
+            "model": model_config.to_dict(),
+            "world_size": world_size,
+            "device": str(device),
+        }
+    )
+
+    if global_rank == 0:
+        if not args.unset_wandb:
+            wandb.config.update(run_config, allow_val_change=True)
+            wandb.save(os.path.abspath(__file__), policy="now")  # save current script
+        # fix tqdm visual length to 80 so that the progress bar
+        # doesn't jump around when changing from external display to laptop
+        pbar = tqdm(total=args.num_training_steps - update_step, desc="Update steps", ncols=80)
+
+    ############ Initialize optimization ############
+    if "galore" in args.optimizer.lower():
+        # make parameters with "rank" to a single group, if param_name has "mlp" or "attn"
+        lowrank_params = []
+        target_modules_list = ["attn", "mlp"]
+        for module_name, module in model.named_modules():
+            if not (isinstance(module, nn.Linear) or isinstance(module, QScaleLinear) or isinstance(module, QLinear)):
+                continue
+            if not any(target_key in module_name for target_key in target_modules_list):
+                continue
+            logger.info(f"Adding {module_name} to GaLore parameters")
+            lowrank_params.append(module.weight)
+
+        id_lowrank_params = [id(p) for p in lowrank_params]
+        # make parameters without "rank" to another group
+        regular_params = [p for p in model.parameters() if id(p) not in id_lowrank_params]
+        # then call low rank optimizer
+        param_groups = [
+            {"params": regular_params},
+            {
+                "params": lowrank_params,
+                "rank": args.rank,
+                "update_proj_gap": args.update_proj_gap,
+                "scale": args.galore_scale,
+                "proj_type": args.proj_type,
+                "quant": args.proj_quant,
+                "quant_n_bit": args.proj_bits,
+                "quant_group_size": args.proj_group_size,
+                "cos_threshold": args.cos_threshold,
+                "gamma_proj": args.gamma_proj,
+                "queue_size": args.queue_size,
+            },
+        ]
+    elif "apollo" in args.optimizer.lower():
+        # make parameters with "rank" to a single group, if param_name has "mlp" or "attn"
+        lowrank_params = []
+        target_modules_list = ["attn", "mlp"]
+        for module_name, module in model.named_modules():
+            if not (isinstance(module, nn.Linear) or isinstance(module, QScaleLinear) or isinstance(module, QLinear)):
+                continue
+            if not any(target_key in module_name for target_key in target_modules_list):
+                continue
+            logger.info(f"Adding {module_name} to APOLLO parameters")
+            lowrank_params.append(module.weight)
+
+        id_lowrank_params = [id(p) for p in lowrank_params]
+        # make parameters without "rank" to another group
+        regular_params = [p for p in model.parameters() if id(p) not in id_lowrank_params]
+        # then call low rank optimizer
+        param_groups = [
+            {"params": regular_params},
+            {
+                "params": lowrank_params,
+                "rank": args.rank,
+                "update_proj_gap": args.update_proj_gap,
+                "scale": args.apollo_scale,
+                "proj_type": args.proj_type,
+                "proj": args.proj,
+                "scale_type": args.scale_type,
+            },
+        ]
+    elif "conda" in args.optimizer.lower():
+        # make parameters with "rank" to a single group, if param_name has "mlp" or "attn"
+        lowrank_params = []
+        target_modules_list = ["attn", "mlp"]
+        for module_name, module in model.named_modules():
+            if not (isinstance(module, nn.Linear) or isinstance(module, QScaleLinear) or isinstance(module, QLinear)):
+                continue
+            if not any(target_key in module_name for target_key in target_modules_list):
+                continue
+            logger.info(f"Adding {module_name} to conda parameters")
+            lowrank_params.append(module.weight)
+
+        id_lowrank_params = [id(p) for p in lowrank_params]
+        # make parameters without "rank" to another group
+        regular_params = [p for p in model.parameters() if id(p) not in id_lowrank_params]
+        # then call low rank optimizer
+        param_groups = [
+            {"params": regular_params},
+            {
+                "params": lowrank_params,
+                "rank": args.rank,
+                "update_proj_gap": args.update_proj_gap,
+                "scale": args.apollo_scale,
+                "proj_type": args.proj_type,
+                "proj": args.proj,
+                "scale_type": args.scale_type,
+            },
+        ]
+    else:
+        param_groups = None
+        id_lowrank_params = None
+
+    # print params and trainable params
+    logger.info(f"\n{model}\n")
+    logger.info(f"Total params: {sum(p.numel() for p in model.parameters()) / 1_000_000:.2f}M")
+
+    if args.simulation:
+        num_train_params = sum(p.numel() for p in trainable_params)
+    else:
+        num_train_params = sum(p.numel() for p in trainable_params) + sum(p.numel() for p in trainable_params_int8)
+
+    logger.info(f"Trainable params: {num_train_params / 1_000_000:.2f}M")
+    if "q_galore" in args.optimizer.lower():
+        logger.info(
+            f"Trainable params with Q-GaLore enabled: {sum(p.numel() for p in trainable_params_int8) / 1_000_000:.2f}M"
+        )
+    elif "galore" in args.optimizer.lower():
+        logger.info(f"Total params with GaLore enabled: {sum(p.numel() for p in lowrank_params) / 1_000_000:.2f}M")
+    elif "q_apollo" in args.optimizer.lower():
+        logger.info(
+            f"Trainable params with Q-APOLLO enabled: {sum(p.numel() for p in trainable_params_int8) / 1_000_000:.2f}M"
+        )
+    elif "apollo" in args.optimizer.lower():
+        logger.info(f"Total params with APOLLO enabled: {sum(p.numel() for p in lowrank_params) / 1_000_000:.2f}M")
+
+    logger.info(f"Saving model to {args.save_dir} every {args.save_every} update steps")
+
+    model, optimizer, scheduler, layer_wise_flag = setup_optimization(
+        args, model, trainable_params, param_groups, id_lowrank_params, model_config
+    )
+
+    if layer_wise_flag:
+        # will pass optimizer_dict and scheduler_dict out instead of optimizer and scheduler
+        optimizer_dict = optimizer
+        scheduler_dict = scheduler
+
+    # Bug-3 fix: wrap with DDP *before* torch.compile per PyTorch recommendation.
+    # This ensures gradient reduction hooks are correctly installed on the DDP module,
+    # and the compiled graph captures the full DDP+model forward pass.
+    # (Issue-5: optimizer.load_state_dict is called after both DDP and compile below.)
+    if not args.single_gpu:
+        model: LlamaForCausalLM = torch.nn.parallel.DistributedDataParallel(
+            model,
+            device_ids=[local_rank],
+            output_device=local_rank,
+            broadcast_buffers=False,
+        )
+
+    # compile the model (after DDP so the compiled graph includes DDP reduction)
+    if args.compile:
+        print("Compiling the model... (takes a ~minute)")
+        unoptimized_model = model
+
+        # Configure TorchDynamo to suppress errors and fall back to eager mode
+        import torch._dynamo
+        torch._dynamo.config.suppress_errors = args.dynamo_suppress_errors
+        torch._dynamo.config.verbose = False
+        # Set cache size limit to prevent memory issues during long training
+        torch._dynamo.config.cache_size_limit = args.dynamo_cache_limit
+
+        model = torch.compile(model) # requires PyTorch 2.0
+
+    # resume optimizer
+    if args.restore_optimizer and args.continue_from is not None:
+        logger.info("Restoring optimizer and scheduler from the checkpoint")
+        _optimizer_dir = args.continue_from
+        optimizer_checkpoint = torch.load(os.path.join(_optimizer_dir, "optimizer.pt"), map_location="cpu")
+        optimizer.load_state_dict(optimizer_checkpoint["optimizer"])
+        scheduler.load_state_dict(optimizer_checkpoint["scheduler"])
+        update_step = optimizer_checkpoint["update_step"]
+        beginning_step = update_step
+        global_step = optimizer_checkpoint["global_step"]
+        logger.info(f"Optimizer and scheduler restored from {_optimizer_dir}")
+
+    # ##############################
+    # TRAINING LOOP
+    # we use iterable dataset, so we may never go through all the data
+    # ##############################
+    # global steps and others are defined above
+    pad_idx = tokenizer.pad_token_id
+    update_time = time.time()
+    local_step = 0  # when continue_from is used, local_step != global_step
+    total_svd_count = 0
+
+    dataloader_iter = iter(dataloader)
+
+    # Issue-4 fix: accumulate loss across micro-batches so logged loss is the true
+    # gradient-accumulation average, not just the last micro-batch.
+    accumulated_loss = 0.0
+
+    # Skip data if resuming from checkpoint
+    if update_step != 0:
+        skip_batches = args.gradient_accumulation * update_step
+        logger.info(f"Skipping {skip_batches} batches to resume from update step {update_step}")
+        skipped = 0
+        for _ in range(skip_batches):
+            # Issue-6 fix: handle StopIteration during skip so all ranks stay aligned
+            try:
+                next(dataloader_iter)
+            except StopIteration:
+                logger.warning(
+                    f"Dataset exhausted during skip at batch {skipped}/{skip_batches}; "
+                    f"restarting iterator to keep ranks aligned."
+                )
+                dataloader_iter = iter(dataloader)
+                next(dataloader_iter)
+            skipped += 1
+        logger.info(f"Skipped {skipped} batches successfully")
+
+    while update_step <= args.num_training_steps:
+        try:
+            batch = next(dataloader_iter)
+        except StopIteration:
+            logger.info(f"Dataset completed one epoch. Starting new epoch with reshuffled data.")
+            dataloader_iter = iter(dataloader)
+            batch = next(dataloader_iter)
+
+        global_step += 1
+        local_step += 1
+
+        if update_step >= args.num_training_steps:
+            logger.info(f"Reached max number of update steps ({args.num_training_steps}). Stopping training.")
+            logger.info(f"Rank {global_rank} stopping training.")
+            break
+
+        # forward & backward
+        batch = {k: v.to(device) for k, v in batch.items()}
+        labels = batch["input_ids"].clone()
+        labels[labels == pad_idx] = -100
+        tokens_seen += (batch["input_ids"] != pad_idx).sum().item() * world_size
+
+        loss = model(**batch, labels=labels).loss
+
+        scaled_loss = loss / args.gradient_accumulation
+        scaled_loss.backward()
+        accumulated_loss += loss.item()  # Issue-4: accumulate before the continue
+
+        if global_step % args.gradient_accumulation != 0:
+            continue
+
+        # The below code is only executed during the update step
+        # Issue-4: compute average loss over all micro-batches in this accumulation window
+        avg_loss = accumulated_loss / args.gradient_accumulation
+        accumulated_loss = 0.0  # reset for next accumulation window
+        # add grad clipping: TODO: add gradient clipping of int8 weight
+        if args.grad_clipping != 0.0:
+            torch.nn.utils.clip_grad_norm_(trainable_params, args.grad_clipping)
+        # Periodic memory cleanup to prevent symbolic tensor issues during long training
+        if global_step % args.memory_cleanup_frequency == 0:
+            torch.cuda.empty_cache()
+            # Clear TorchDynamo cache to prevent memory accumulation
+            if args.compile:
+                import torch._dynamo
+                torch._dynamo.reset()
+
+        if global_rank == 0:
+            pbar.update(1)
+        if not layer_wise_flag:  # layer-wise updation is done during backward; requires gradient_accumulation equals 1
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+
+        update_step += 1
+        update_time = time.time() - update_time
+
+        # save checkpoint by save_every
+        if local_step > args.gradient_accumulation and update_step % args.save_every == 0 and global_rank == 0:
+            current_model_directory = f"{args.save_dir}/model_{update_step}"
+            logger.info(f"Saving model and optimizer to {current_model_directory}, update step {update_step}")
+            os.makedirs(args.save_dir, exist_ok=True)
+            # Bug-1 fix: unwrap DDP/compiled model for saving; works in both single-GPU and multi-GPU modes
+            unwrapped_model = model.module if hasattr(model, 'module') else model
+            unwrapped_model.save_pretrained(current_model_directory, max_shard_size="500GB", from_pt=True)
+            saving_model_weight(unwrapped_model, f"{current_model_directory}/pytorch_model.bin", args)
+
+            optimizer_checkpoint = {
+                "optimizer": optimizer.state_dict(),
+                "scheduler": scheduler.state_dict(),
+                "update_step": update_step,
+                "global_step": global_step,
+                "config": run_config,
+                "wandb": wandb.run.dir if not args.unset_wandb else None,
+                "dtype": args.dtype,
+            }
+            torch.save(optimizer_checkpoint, f"{current_model_directory}/optimizer.pt")
+
+            training_state_checkpoint = {
+                "global_step": global_step,
+                "update_step": update_step,
+                "tokens_seen": tokens_seen,
+                "tokens_seen_before": tokens_seen_before,
+                "update_time": update_time,
+            }
+            with open(f"{current_model_directory}/training_state.json", "w") as f:
+                json.dump(training_state_checkpoint, f, indent=4)
+
+            # save wandb related info
+            if not args.unset_wandb:
+                wandb_info = {
+                    "wandb_id": wandb.run.id,
+                }
+                with open(f"{args.save_dir}/wandb.json", "w") as f:
+                    json.dump(wandb_info, f, indent=4)
+
+        # evaluation
+        if update_step % args.eval_every == 0:
+            logger.info(f"Performing evaluation at step {update_step}")
+            total_loss, evaluated_on_tokens, perplexity = evaluate_model(
+                model, tokenizer, pad_idx, global_rank, world_size, device, args
+            )
+
+            if global_rank == 0:
+                if not args.unset_wandb:
+                    wandb.log(
+                        {
+                            "eval_loss": total_loss,
+                            "eval_perplexity": perplexity,
+                            "eval_tokens": evaluated_on_tokens,
+                        },
+                        step=update_step,
+                    )
+            logger.info(f"Eval loss at step {update_step}: {total_loss}, Eval perplexity: {perplexity}")
+
+        if not layer_wise_flag:
+            lr = optimizer.param_groups[0]["lr"]
+        else:
+            lr = list(optimizer_dict.values())[0].param_groups[0]["lr"]
+        tokens_in_update = tokens_seen - tokens_seen_before
+        tokens_seen_before = tokens_seen
+        batches_in_update = args.gradient_accumulation * world_size
+        if not layer_wise_flag:
+            total_svd_count = getting_svd_cnt(optimizer)
+        else:
+            total_svd_count = 0
+
+        # Build extra lr metrics for hybrid optimizer
+        hybrid_lr_log = {}
+        if isinstance(optimizer, HybridOptimizer):
+            if optimizer.muon is not None:
+                hybrid_lr_log["lr_muon"]   = optimizer.muon.param_groups[0]["lr"]
+            if optimizer.apollo is not None:
+                hybrid_lr_log["lr_apollo"] = optimizer.apollo.param_groups[0]["lr"]
+            hybrid_lr_log["lr_adamw"] = optimizer.adamw.param_groups[0]["lr"]
+
+        if global_rank == 0:
+            if not args.unset_wandb:
+                log_dict = {
+                    "loss": avg_loss,
+                    "lr": lr,
+                    "update_step": update_step,
+                    "tokens_seen": tokens_seen,
+                    "total_svd_count": total_svd_count,
+                    "throughput_tokens": tokens_in_update / update_time,
+                    "throughput_examples": args.total_batch_size / update_time,
+                    "throughput_batches": batches_in_update / update_time,
+                }
+                log_dict.update(hybrid_lr_log)
+                wandb.log(log_dict, step=update_step)
+        update_time = time.time()
+
+    # ##############################
+    # END of training loop
+    # ##############################
+    logger.info("Training finished")
+    if global_rank == 0:
+        pbar.close()
+
+    current_model_directory = f"{args.save_dir}/model_{update_step}"
+    if global_rank == 0 and not os.path.exists(current_model_directory):
+        logger.info(f"Saving model and optimizer to {current_model_directory}, update step {update_step}")
+        os.makedirs(args.save_dir, exist_ok=True)
+        # Bug-1 fix: unwrap DDP/compiled model for saving; works in both single-GPU and multi-GPU modes
+        unwrapped_model = model.module if hasattr(model, 'module') else model
+        unwrapped_model.save_pretrained(current_model_directory, max_shard_size="500GB", from_pt=True)
+        saving_model_weight(unwrapped_model, f"{current_model_directory}/pytorch_model.bin", args)
+
+        optimizer_checkpoint = {
+            "optimizer": optimizer.state_dict(),
+            "scheduler": scheduler.state_dict(),
+            "update_step": update_step,
+            "global_step": global_step,
+            "config": run_config,
+            "wandb": wandb.run.dir if not args.unset_wandb else None,
+            "dtype": args.dtype,
+        }
+        torch.save(optimizer_checkpoint, f"{current_model_directory}/optimizer.pt")
+
+        training_state_checkpoint = {
+            "global_step": global_step,
+            "update_step": update_step,
+            "tokens_seen": tokens_seen,
+            "tokens_seen_before": tokens_seen_before,
+            "update_time": update_time,
+        }
+        with open(f"{current_model_directory}/training_state.json", "w") as f:
+            json.dump(training_state_checkpoint, f, indent=4)
+
+    # Final evaluation
+    logger.info("Running final evaluation")
+    model.eval()
+    del loss, optimizer, scheduler
+    import gc
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    total_loss, evaluated_on_tokens, perplexity = evaluate_model(model, tokenizer, pad_idx, global_rank, world_size, device, args)
+
+    if global_rank == 0:
+        if not args.unset_wandb:
+            wandb.log(
+                {
+                    "final_eval_loss": total_loss,
+                    "final_eval_perplexity": perplexity,
+                    "final_eval_tokens": evaluated_on_tokens,
+                },
+                step=update_step,
+            )
+        logger.info(f"Final eval loss: {total_loss}, Final eval perplexity: {perplexity}")
+
+    logger.info("Script finished successfully")
+    print(f"Rank {global_rank} finished successfully")
+
+
+if __name__ == "__main__":
+    print("Starting script")
+    args = parse_args(None)
+    main(args)

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/files/requirements.txt

@@ -0,0 +1,142 @@
+sac==0.1.0
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+torchaudio==2.11.0
+nvidia-cusparselt-cu12==0.7.1
+mpmath==1.3.0
+typing_extensions==4.15.0
+triton==3.4.0
+sympy==1.14.0
+pillow==12.2.0
+nvidia-nvtx-cu12==12.8.90
+nvidia-nvjitlink-cu12==12.8.93
+nvidia-nccl-cu12==2.27.3
+nvidia-curand-cu12==10.3.9.90
+nvidia-cufile-cu12==1.13.1.3
+nvidia-cuda-runtime-cu12==12.8.90
+nvidia-cuda-nvrtc-cu12==12.8.93
+nvidia-cuda-cupti-cu12==12.8.90
+nvidia-cublas-cu12==12.8.4.1
+numpy==2.2.6
+networkx==3.4.2
+MarkupSafe==3.0.3
+aiohappyeyeballs==2.6.1
+filelock==3.28.0
+nvidia-cusparse-cu12==12.5.8.93
+nvidia-cufft-cu12==11.3.3.83
+nvidia-cudnn-cu12==9.10.2.21
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.3.90
+torch==2.8.0+cu128
+torchvision==0.23.0+cu128
+pytz==2026.1.post1
+xxhash==3.6.0
+urllib3==2.6.3
+tzdata==2026.1
+tqdm==4.67.3
+six==1.17.0
+safetensors==0.7.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyarrow==23.0.1
+psutil==7.2.2
+propcache==0.4.1
+multidict==6.7.1
+idna==3.11
+hf-xet==1.4.3
+h11==0.16.0
+fsspec==2026.2.0
+frozenlist==1.8.0
+exceptiongroup==1.3.1
+dill==0.4.1
+charset-normalizer==3.4.7
+certifi==2026.2.25
+attrs==26.1.0
+async-timeout==5.0.1
+yarl==1.23.0
+requests==2.33.1
+python-dateutil==2.9.0.post0
+multiprocess==0.70.19
+httpcore==1.0.9
+anyio==4.13.0
+aiosignal==1.4.0
+pandas==2.3.3
+huggingface_hub==0.36.2
+httpx==0.28.1
+aiohttp==3.13.5
+tokenizers==0.22.2
+accelerate==1.13.0
+transformers==4.57.3
+datasets==4.8.4
+peft==0.19.1
+pytorch-ranger==0.1.1
+lion-pytorch==0.2.4
+bitsandbytes==0.49.2
+torch-optimizer==0.3.0
+apollo-torch==1.0.3
+nvidia-ml-py==13.590.48
+typing-inspection==0.4.2
+threadpoolctl==3.6.0
+smmap==5.0.3
+sentry-sdk==2.58.0
+scipy==1.15.3
+pyparsing==3.3.2
+pydantic_core==2.46.3
+protobuf==7.34.1
+platformdirs==4.9.6
+nvitop==1.6.2
+loguru==0.7.3
+kiwisolver==1.5.0
+joblib==1.5.3
+fonttools==4.62.1
+cycler==0.12.1
+contourpy==1.3.2
+click==8.3.2
+annotated-types==0.7.0
+scikit-learn==1.7.2
+pydantic==2.13.3
+modelscope==1.35.4
+matplotlib==3.10.8
+gitdb==4.0.12
+seaborn==0.13.2
+GitPython==3.1.47
+wandb==0.26.0
+sac==0.1.0
+nvidia-ml-py3==7.352.0
+gitignore_parser==0.1.13
+durationpy==0.10
+dotmap==1.3.30
+wrapt==2.1.2
+websocket-client==1.9.0
+typeguard==4.5.1
+tabulate==0.9.0
+pycparser==3.0
+pyasn1==0.6.3
+py==1.11.0
+oauthlib==3.3.1
+jmespath==1.1.0
+invoke==3.0.3
+elasticsearch==7.17.13
+docutils==0.19
+decorator==5.2.1
+confluent-kafka==2.14.0
+colorama==0.4.6
+bcrypt==5.0.0
+rsa==4.7.2
+retry==0.9.2
+requests-oauthlib==2.0.0
+Deprecated==1.3.1
+cffi==2.0.0
+botocore==1.42.92
+s3transfer==0.16.0
+PyNaCl==1.6.2
+kubernetes==35.0.0
+cryptography==46.0.7
+paramiko==4.0.0
+boto3==1.42.92
+awscli==1.44.82
+megfile==2.2.10.post1
+refile==7.2.7.post3
+brainpp==2.7.12.16

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/logs/debug-internal.log

@@ -0,0 +1,15 @@
+{"time":"2026-04-23T20:24:22.4996429+08:00","level":"INFO","msg":"wandb-core"}
+{"time":"2026-04-23T20:24:22.499901972+08:00","level":"INFO","msg":"stream: starting","core version":"0.26.0"}
+{"time":"2026-04-23T20:24:22.63528813+08:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
+{"time":"2026-04-23T20:24:22.635312622+08:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
+{"time":"2026-04-23T20:24:22.635332006+08:00","level":"INFO","msg":"stream: created new stream","id":"mawza3ul"}
+{"time":"2026-04-23T20:24:22.635539387+08:00","level":"INFO","msg":"handler: started"}
+{"time":"2026-04-23T20:24:22.636204984+08:00","level":"INFO","msg":"stream: started"}
+{"time":"2026-04-23T20:24:22.636292599+08:00","level":"INFO","msg":"writer: started","stream_id":"mawza3ul"}
+{"time":"2026-04-23T20:24:22.636304487+08:00","level":"INFO","msg":"sender: started"}
+{"time":"2026-04-23T20:24:22.637327692+08:00","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
+{"time":"2026-04-23T20:24:22.637343161+08:00","level":"WARN","msg":"runupserter: server does not expand metric globs but the x_server_side_expand_glob_metrics setting is set; ignoring"}
+{"time":"2026-04-24T01:14:24.730723575+08:00","level":"INFO","msg":"stream: closing"}
+{"time":"2026-04-24T01:14:24.753028533+08:00","level":"INFO","msg":"handler: closed"}
+{"time":"2026-04-24T01:14:24.753387495+08:00","level":"INFO","msg":"sender: closed"}
+{"time":"2026-04-24T01:14:24.753400977+08:00","level":"INFO","msg":"stream: closed"}

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/logs/debug.log

@@ -0,0 +1,22 @@
+2026-04-23 20:24:22,097 INFO    MainThread:342 [wandb_setup.py:_flush():81] Current SDK version is 0.26.0
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_setup.py:_flush():81] Configure stats pid to 342
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:setup_run_log_directory():721] Logging user logs to exp_remain_h200/work_dirs/350m/train_350m_conda_lr1e_2_scale0_25_rank256_gap2000_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/logs/debug.log
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:setup_run_log_directory():722] Logging internal logs to exp_remain_h200/work_dirs/350m/train_350m_conda_lr1e_2_scale0_25_rank256_gap2000_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/logs/debug-internal.log
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:init():848] calling init triggers
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:init():853] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-04-23 20:24:22,098 INFO    MainThread:342 [wandb_init.py:init():896] starting backend
+2026-04-23 20:24:22,494 INFO    MainThread:342 [wandb_init.py:init():911] sending inform_init request
+2026-04-23 20:24:22,498 INFO    MainThread:342 [wandb_init.py:init():919] backend started and connected
+2026-04-23 20:24:22,498 INFO    MainThread:342 [wandb_init.py:init():989] updated telemetry
+2026-04-23 20:24:22,526 INFO    MainThread:342 [wandb_init.py:init():1013] communicating run to backend with 90.0 second timeout
+2026-04-23 20:24:22,638 INFO    MainThread:342 [wandb_init.py:init():1058] starting run threads in backend
+2026-04-23 20:24:22,712 INFO    MainThread:342 [wandb_run.py:_console_start():2542] atexit reg
+2026-04-23 20:24:22,712 INFO    MainThread:342 [wandb_run.py:_redirect():2391] redirect: wrap_raw
+2026-04-23 20:24:22,712 INFO    MainThread:342 [wandb_run.py:_redirect():2460] Wrapping output streams.
+2026-04-23 20:24:22,712 INFO    MainThread:342 [wandb_run.py:_redirect():2483] Redirects installed.
+2026-04-23 20:24:22,714 INFO    MainThread:342 [wandb_init.py:init():1098] run started, returning control to user process
+2026-04-23 20:24:49,867 INFO    MainThread:342 [wandb_run.py:_config_callback():1403] config_cb None None {'model_config': 'configs/llama_350m.json', 'exp_config': 'exp_v2/configs/llama_350m.json', 'eval_every': 1000, 'save_every': 60000, 'dtype': 'bfloat16', 'seed': 0, 'compile': True, 'dynamo_suppress_errors': True, 'dynamo_cache_limit': 10000, 'memory_cleanup_frequency': 10000, 'resume_step': None, 'restore_optimizer': False, 'continue_from': None, 'single_gpu': False, 'save_dir': 'exp_remain_h200/work_dirs/350m/train_350m_conda_lr1e_2_scale0_25_rank256_gap2000_20260423_202415', 'use_hf_model': False, 'workers': 12, 'batch_size': 128, 'gradient_accumulation': 1, 'total_batch_size': 512, 'warmup_steps': 6000, 'num_training_steps': 60000, 'max_train_tokens': None, 'optimizer': 'conda', 'max_length': 256, 'scheduler': 'cosine', 'min_lr_ratio': 0.1, 'weight_decay': 0.0, 'grad_clipping': 0.0, 'activation_checkpointing': False, 'data_path': '/mnt/shared-storage-gpfs2/finebio-shared/optimizer/dataset/C4/en', 'data_name': 'en', 'tags': None, 'name': 'test', 'project': 'test', 'unset_wandb': False, 'entity': None, 'wandb_dir': 'exp_remain_h200/work_dirs/350m/train_350m_conda_lr1e_2_scale0_25_rank256_gap2000_20260423_202415', 'beta1': 0.9, 'beta2': 0.99, 'beta3': 0.99, 'eps': 1e-08, 'rank': 256, 'update_proj_gap': 2000, 'galore_scale': 1.0, 'proj_type': 'std', 'proj_quant': False, 'proj_bits': 8, 'proj_group_size': 256, 'weight_quant': False, 'weight_bits': 8, 'weight_group_size': 256, 'stochastic_round': False, 'simulation': False, 'cos_threshold': 1, 'gamma_proj': 2, 'queue_size': 5, 'proj': 'random', 'scale_type': 'channel', 'apollo_scale': 0.25, 'scale_front': False, 'n_clusters': 3, 'scale_update_freq': 500, 'scale_level': '1,0,1,1', 'scale_bound': None, 'metric': 'mean', 'align_grad': False, 'dim': 4096, 'n_heads': 32, 'muon_ns_steps': 5, 'muon_momentum': 0.95, 'nproc_per_node': 4, 'max_lr': 0.01, 'total_params_M': 367.96928, 'dataset': 'c4', 'model': {'vocab_size': 32000, 'max_position_embeddings': 2048, 'hidden_size': 1024, 'intermediate_size': 2736, 'num_hidden_layers': 24, 'num_attention_heads': 16, 'num_key_value_heads': 16, 'hidden_act': 'silu', 'initializer_range': 0.02, 'rms_norm_eps': 1e-06, 'pretraining_tp': 1, 'use_cache': True, 'rope_theta': 10000.0, 'rope_scaling': None, 'attention_bias': False, 'attention_dropout': 0.0, 'mlp_bias': False, 'head_dim': 64, 'return_dict': True, 'output_hidden_states': False, 'torchscript': False, 'dtype': None, 'pruned_heads': {}, 'tie_word_embeddings': False, 'chunk_size_feed_forward': 0, 'is_encoder_decoder': False, 'is_decoder': False, 'cross_attention_hidden_size': None, 'add_cross_attention': False, 'tie_encoder_decoder': False, 'architectures': ['LLaMAForCausalLM'], 'finetuning_task': None, 'id2label': {0: 'LABEL_0', 1: 'LABEL_1'}, 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, 'task_specific_params': None, 'problem_type': None, 'tokenizer_class': None, 'prefix': None, 'bos_token_id': 0, 'pad_token_id': -1, 'eos_token_id': 1, 'sep_token_id': None, 'decoder_start_token_id': None, 'max_length': 20, 'min_length': 0, 'do_sample': False, 'early_stopping': False, 'num_beams': 1, 'temperature': 1.0, 'top_k': 50, 'top_p': 1.0, 'typical_p': 1.0, 'repetition_penalty': 1.0, 'length_penalty': 1.0, 'no_repeat_ngram_size': 0, 'encoder_no_repeat_ngram_size': 0, 'bad_words_ids': None, 'num_return_sequences': 1, 'output_scores': False, 'return_dict_in_generate': False, 'forced_bos_token_id': None, 'forced_eos_token_id': None, 'remove_invalid_values': False, 'exponential_decay_length_penalty': None, 'suppress_tokens': None, 'begin_suppress_tokens': None, 'num_beam_groups': 1, 'diversity_penalty': 0.0, '_name_or_path': 'configs/llama_350m.json', 'transformers_version': '4.57.3', 'max_sequence_length': 1024, 'model_type': 'llama', 'tf_legacy_loss': False, 'use_bfloat16': False, 'output_attentions': False}, 'world_size': 4, 'device': 'cuda:0'}
+2026-04-24 01:14:24,727 INFO    wandb-AsyncioManager-main:342 [service_client.py:_forward_responses():134] Reached EOF.
+2026-04-24 01:14:24,730 INFO    wandb-AsyncioManager-main:342 [mailbox.py:close():155] Closing mailbox, abandoning 0 handles.

350m/conda_lr1e_2_b1_0_9_b2_0_99_eps_1e_8_scale_0_25_rank256_T2000_H200_ppl16_4542_20260423_202415/wandb/offline-run-20260423_202422-mawza3ul/run-mawza3ul.wandb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:962d8f0f208560ca403ae8a08e4f5b16a4672f36756493a2b4a1f93e72e37218
+size 59876049