Want learn with model and optim config variable as float16 but i can't
State: closed | Pull request: no

Hello. 
I'm beginner of using NeMo framework.

Following the tutorial_pretraining.py (https://docs.nvidia.com/nemo-framework/user-guide/latest/nemo-2.0/quickstart.html)
I want to make small model (mamba7_780m) be learned with my old gpus ( Geforce RTX 2080 Ti -12GB * 2) .
So i wanted to chang model, optimizer variables like bellow. but bf16, fp16 variables are not changed. 
The other variable are changed well (ex - adam_beta1) . 

  #Add
  recipe.model.config.fp16 = True
  recipe.model.config.bf16 = False
  recipe.**optim.config.fp16 = True**
  recipe.**optim.config.bf16 = False**
  recipe.optim.config.adam_beta1=0.123
  #Error
   [NeMo I 2025-01-03 05:50:17 utils:302] Setting up optimizer with config OptimizerConfig(optimizer='adam', lr=0.0003, min_lr=None, decoupled_lr=None, decoupled_min_lr=None, weight_decay=0.1, **fp16=False, bf16=True**, params_dtype=torch.bfloat16, loss_scale=None, initial_loss_scale=4294967296, min_loss_scale=1.0, loss_scale_window=1000, hysteresis=2, adam_beta1=0.123, adam_beta2=0.95, adam_eps=1e-05, sgd_momentum=0.9, use_distributed_optimizer=False, overlap_param_gather_with_optimizer_step=False, clip_grad=1.0, log_num_zeros_in_grad=False, barrier_with_L1_time=False, timers=None, config_logger_dir='')

How can i solve it?

The code i edited
```
def configure_recipe(nodes: int = 1, gpus_per_node: int = 1):
    recipe = llm.mamba2_780m.pretrain_recipe(
        dir="/checkpoints/nemotron", # Path to store checkpoints
        name="mamba_pretraining",
        num_nodes=nodes,
        num_gpus_per_node=gpus_per_node,
        tensor_parallelism = 1
        
    )
    
    recipe.trainer.val_check_interval = 100
    recipe.data.global_batch_size = 2
    recipe.data.micro_batch_size = 1
    recipe.data.seq_length = 256
    # recipe.model.config.num_layers = 24
    recipe.model.config.num_attention_heads = 2
    recipe.model.config.fp16 = True
    recipe.model.config.bf16 = False
    recipe.optim.config.fp16 = True
    recipe.optim.config.bf16 = False
    recipe.optim.config.use_distributed_optimizer = False
    recipe.optim.config.params_dtype = torch.float16
    recipe.optim.config.adam_beta1=0.123

    return recipe

def local_executor_torchrun(nodes: int = 1, devices: int = 1) -> run.LocalExecutor:
    # Env vars for jobs are configured here
    env_vars = {
        "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
        "NCCL_NVLS_ENABLE": "0",
        "NVTE_DP_AMAX_REDUCE_INTERVAL": "0",
        "NVTE_ASYNC_AMAX_REDUCTION": "1",
    }

    executor = run.LocalExecutor(ntasks_per_node=devices, launcher="torchrun", env_vars=env_vars)

    return executor

def run_pretraining():
    recipe = configure_recipe()
    executor = local_executor_torchrun(nodes=recipe.trainer.num_nodes, devices=recipe.trainer.devices)

    run.run(recipe, executor=executor, name="mamba2_780m_pretraining")

# This condition is necessary for the script to be compatible with Python's multiprocessing module.
if __name__ == "__main__":
    run_pretraining()
```
Here is the result of code
```

──────────────────────────────── Entering Experiment mamba2_780m_pretraining with id: mamba2_780m_pretraining_1735883405 ─────────────────────────────────
[05:50:05] INFO     Log directory is:                                                                                               local_scheduler.py:777
                    /root/.nemo_run/experiments/mamba2_780m_pretraining/mamba2_780m_pretraining_1735883405/mamba2_780m_pretraining                        
[05:50:05] Launching job mamba2_780m_pretraining for experiment mamba2_780m_pretraining                                                  experiment.py:724
           INFO     Log directory is:                                                                                               local_scheduler.py:777
                    /root/.nemo_run/experiments/mamba2_780m_pretraining/mamba2_780m_pretraining_1735883405/mamba2_780m_pretraining                        
           INFO     Launched app: local_persistent://nemo_run/mamba2_780m_pretraining-ht7tmrvqhgqrgd                                       launcher.py:111
           INFO     AppStatus:                                                                                                             launcher.py:114
                        State: RUNNING                                                                                                                    
                        Num Restarts: 0                                                                                                                   
                        Roles:                                                                                                                            
                        Msg: <NONE>                                                                                                                       
                        Structured Error Msg: <NONE>                                                                                                      
                        UI URL:                                                                                                                           
                    file:///root/.nemo_run/experiments/mamba2_780m_pretraining/mamba2_780m_pretraining_1735883405/mamba2_780m_pretraining/                
                    nemo_run/mamba2_780m_pretraining-ht7tmrvqhgqrgd                                                                                       
                                                                                                                                                          
────────────────────────────────────────── Waiting for Experiment mamba2_780m_pretraining_1735883405 to finish ───────────────────────────────────────────

Experiment Status for mamba2_780m_pretraining_1735883405

Task 0: mamba2_780m_pretraining
- Status: RUNNING
- Executor: LocalExecutor
- Job id: mamba2_780m_pretraining-ht7tmrvqhgqrgd
- Local Directory: /root/.nemo_run/experiments/mamba2_780m_pretraining/mamba2_780m_pretraining_1735883405/mamba2_780m_pretraining

           INFO     Waiting for job mamba2_780m_pretraining-ht7tmrvqhgqrgd to finish [log=True]...                                         launcher.py:132
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194] Starting elastic_operator with launch configs:
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   entrypoint       : nemo_run.core.runners.fdl_runner
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   min_nodes        : 1
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   max_nodes        : 1
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   nproc_per_node   : 1
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   run_id           : 868
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   rdzv_backend     : c10d
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   rdzv_endpoint    : localhost:0
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   rdzv_configs     : {'timeout': 900}
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   max_restarts     : 0
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   monitor_interval : 0.1
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   log_dir          : /root/.nemo_run/experiments/mamba2_780m_pretraining/mamba2_780m_pretraining_1735883405/mamba2_780m_pretraining/nemo_run/mamba2_780m_pretraining-ht7tmrvqhgqrgd/torchelastic/mamba2_780m_pretraining
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194]   metrics_cfg      : {}
retraining/0 I0103 05:50:06.640000 17575 torch/distributed/launcher/api.py:194] 
retraining/0 I0103 05:50:06.642000 17575 torch/distributed/elastic/agent/server/api.py:845] [default] starting workers for entrypoint: python
retraining/0 I0103 05:50:06.642000 17575 torch/distributed/elastic/agent/server/api.py:662] [default] Rendezvous'ing worker group
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525] [default] Rendezvous complete for workers. Result:
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   restart_count=0
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   master_addr=localhost
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   master_port=36739
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   group_rank=0
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   group_world_size=1
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   local_ranks=[0]
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   role_ranks=[0]
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   global_ranks=[0]
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   role_world_sizes=[1]
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525]   global_world_sizes=[1]
retraining/0 I0103 05:50:06.814000 17575 torch/distributed/elastic/agent/server/api.py:525] 
retraining/0 I0103 05:50:06.815000 17575 torch/distributed/elastic/agent/server/api.py:670] [default] Starting worker group
retraining/0 I0103 05:50:06.815000 17575 torch/distributed/elastic/agent/server/local_elastic_agent.py:291] use_agent_store: True
retraining/0 I0103 05:50:06.815000 17575 torch/distributed/elastic/agent/server/local_elastic_agent.py:192] Environment variable 'TORCHELASTIC_ENABLE_FILE_TIMER' not found. Do not start FileTimerServer.
retraining/0 I0103 05:50:06.815000 17575 torch/distributed/elastic/agent/server/local_elastic_agent.py:229] Environment variable 'TORCHELASTIC_HEALTH_CHECK_PORT' not found. Do not start health check.
retraining/0 [default0]:[NeMo W 2025-01-03 05:50:14 nemo_logging:361] /usr/local/lib/python3.10/dist-packages/pyannote/core/notebook.py:134: MatplotlibDeprecationWarning: The get_cmap function was deprecated in Matplotlib 3.7 and will be removed in 3.11. Use ``matplotlib.colormaps[name]`` or ``matplotlib.colormaps.get_cmap()`` or ``pyplot.get_cmap()`` instead.
retraining/0 [default0]:      cm = get_cmap("Set1")
retraining/0 [default0]:    
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 tokenizer_utils:185] Getting HuggingFace AutoTokenizer with pretrained_model_name: EleutherAI/gpt-neox-20b
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 tokenizer_utils:185] Getting HuggingFace AutoTokenizer with pretrained_model_name: EleutherAI/gpt-neox-20b
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 nemo_logger:145] Experiments will be logged at /checkpoints/nemotron/mamba_pretraining
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:426] Rank 0 has data parallel group : [0]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:432] Rank 0 has combined group of data parallel and context parallel : [0]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:437] All data parallel group ranks with context parallel combined: [[0]]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:440] Ranks 0 has data parallel rank: 0
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:448] Rank 0 has context parallel group: [0]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:451] All context parallel group ranks: [[0]]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:452] Ranks 0 has context parallel rank: 0
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:459] Rank 0 has model parallel group: [0]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:460] All model parallel group ranks: [[0]]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:469] Rank 0 has tensor model parallel group: [0]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:473] All tensor model parallel group ranks: [[0]]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:474] Rank 0 has tensor model parallel rank: 0
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:494] Rank 0 has pipeline model parallel group: [0]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:506] Rank 0 has embedding group: [0]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:512] All pipeline model parallel group ranks: [[0]]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:513] Rank 0 has pipeline model parallel rank 0
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:514] All embedding group ranks: [[0]]
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:16 megatron_init:515] Rank 0 has embedding rank: 0
retraining/0 [default0]:GPU available: True (cuda), used: True
retraining/0 [default0]:TPU available: False, using: 0 TPU cores
retraining/0 [default0]:HPU available: False, using: 0 HPUs
retraining/0 [default0]:[NeMo W 2025-01-03 05:50:16 nemo_logger:123] No version folders would be created under the log folder as 'resume_if_exists' is enabled.
retraining/0 [default0]:[NeMo W 2025-01-03 05:50:16 nemo_logger:173] "update_logger_directory" is True. Overwriting tensorboard logger "save_dir" to /checkpoints/nemotron/tb_logs
retraining/0 [default0]:[NeMo W 2025-01-03 05:50:16 nemo_logger:189] The Trainer already contains a ModelCheckpoint callback. This will be overwritten.
retraining/0 [default0]:[NeMo W 2025-01-03 05:50:16 nemo_logger:212] The checkpoint callback was told to monitor a validation value and trainer's max_steps was set to 100. Please ensure that max_steps will run for at least 1 epochs to ensure that checkpointing will not error out.
retraining/0 [default0]:[NeMo W 2025-01-03 05:50:16 resume:228] There were no checkpoints found in checkpoint_dir or no checkpoint folder at checkpoint_dir :/checkpoints/nemotron/mamba_pretraining/checkpoints. Training from scratch.
retraining/0 [default0]:Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/1
retraining/0 [default0]:----------------------------------------------------------------------------------------------------
retraining/0 [default0]:distributed_backend=nccl
retraining/0 [default0]:All distributed processes registered. Starting with 1 processes
retraining/0 [default0]:----------------------------------------------------------------------------------------------------
retraining/0 [default0]:
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 base:44] Padded vocab_size: 50288, original vocab_size: 50277, dummy tokens: 11.
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] Using hybrid override pattern
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] The override pattern matches the overridden pattern
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] Hybrid allocation (M is mamba, * is attention, - is mlp):
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] MMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMM
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] 0 attention layers in 48 total layers.
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] Target attention ratio: 0.00. Actual attention ratio: 0.00.
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] 0 mlp layers in 48 total layers.
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] Target mlp ratio: 0.00. Actual mlp ratio: 0.00.
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 megatron_strategy:327] Copying Trainer's 'max_steps' (100) to LR scheduler's 'max_steps'.
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 num_microbatches_calculator:228] setting number of microbatches to constant 2
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 megatron_parallel:549]  > number of parameters on (tensor, pipeline) model parallel rank (0, 0): 857403648
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] Setting up DistributedDataParallel with config DistributedDataParallelConfig(grad_reduce_in_fp32=True, overlap_grad_reduce=True, overlap_param_gather=True, align_param_gather=False, use_distributed_optimizer=False, num_distributed_optimizer_instances=1, check_for_nan_in_grad=True, bucket_size=40000000, average_in_collective=False, fp8_param_gather=False)
retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:323] Number of buckets for gradient all-reduce / reduce-scatter: 18

[...]

retraining/0 [default0]:[NeMo I 2025-01-03 05:50:17 utils:302] Setting up optimizer with config OptimizerConfig(optimizer='adam', lr=0.0003, min_lr=None, decoupled_lr=None, decoupled_min_lr=None, weight_decay=0.1, fp16=False, bf16=True, params_dtype=torch.bfloat16, loss_scale=None, initial_loss_scale=4294967296, min_loss_scale=1.0, loss_scale_window=1000, hysteresis=2, adam_beta1=0.123, adam_beta2=0.95, adam_eps=1e-05, sgd_momentum=0.9, use_distributed_optimizer=False, overlap_param_gather_with_optimizer_step=False, clip_grad=1.0, log_num_zeros_in_grad=False, barrier_with_L1_time=False, timers=None, config_logger_dir='')
retraining/0 [default0]:LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]
retraining/0 [default0]:
retraining/0 [default0]:  | Name   | Type | Params | Mode 
retraining/0 [default0]:----------------------------------------
retraining/0 [default0]:0 | module | DDP  | 857 M  | train
retraining/0 [default0]:----------------------------------------
retraining/0 [default0]:857 M     Trainable params
retraining/0 [default0]:0         Non-trainable params
retraining/0 [default0]:857 M     Total params
retraining/0 [default0]:3,429.615 Total estimated model params size (MB)
retraining/0 [default0]:394       Modules in train mode
retraining/0 [default0]:0         Modules in eval mode
retraining/0 [default0]:[rank0]: Traceback (most recent call last):
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/compiler.py", line 292, in make_cubin
retraining/0 [default0]:[rank0]:     subprocess.run(cmd, shell=True, check=True)
retraining/0 [default0]:[rank0]:   File "/usr/lib/python3.10/subprocess.py", line 526, in run
retraining/0 [default0]:[rank0]:     raise CalledProcessError(retcode, process.args,
retraining/0 [default0]:[rank0]: subprocess.CalledProcessError: Command '/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/bin/ptxas -lineinfo -v --gpu-name=sm_75 /tmp/tmppa4c09n_.ptx -o /tmp/tmppa4c09n_.ptx.o 2> /tmp/tmp39of6etg.log' returned non-zero exit status 255.
retraining/0 [default0]:
retraining/0 [default0]:[rank0]: During handling of the above exception, another exception occurred:
retraining/0 [default0]:
retraining/0 [default0]:[rank0]: Traceback (most recent call last):
retraining/0 [default0]:[rank0]:   File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
retraining/0 [default0]:[rank0]:     return _run_code(code, main_globals, None,
retraining/0 [default0]:[rank0]:   File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
retraining/0 [default0]:[rank0]:     exec(code, run_globals)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo-Run/src/nemo_run/core/runners/fdl_runner.py", line 66, in <module>
retraining/0 [default0]:[rank0]:     fdl_runner_app()
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/main.py", line 338, in __call__
retraining/0 [default0]:[rank0]:     raise e
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/main.py", line 321, in __call__
retraining/0 [default0]:[rank0]:     return get_command(self)(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1157, in __call__
retraining/0 [default0]:[rank0]:     return self.main(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/core.py", line 665, in main
retraining/0 [default0]:[rank0]:     return _main(
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/core.py", line 197, in _main
retraining/0 [default0]:[rank0]:     rv = self.invoke(ctx)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 1434, in invoke
retraining/0 [default0]:[rank0]:     return ctx.invoke(self.callback, **ctx.params)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/click/core.py", line 783, in invoke
retraining/0 [default0]:[rank0]:     return __callback(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/typer/main.py", line 703, in wrapper
retraining/0 [default0]:[rank0]:     return callback(**use_params)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo-Run/src/nemo_run/core/runners/fdl_runner.py", line 62, in fdl_direct_run
retraining/0 [default0]:[rank0]:     fdl_fn()
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/collections/llm/api.py", line 150, in pretrain
retraining/0 [default0]:[rank0]:     return train(
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/collections/llm/api.py", line 107, in train
retraining/0 [default0]:[rank0]:     trainer.fit(model, data)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py", line 538, in fit
retraining/0 [default0]:[rank0]:     call._call_and_handle_interrupt(
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
retraining/0 [default0]:[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
retraining/0 [default0]:[rank0]:     return function(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py", line 574, in _fit_impl
retraining/0 [default0]:[rank0]:     self._run(model, ckpt_path=ckpt_path)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py", line 981, in _run
retraining/0 [default0]:[rank0]:     results = self._run_stage()
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py", line 1023, in _run_stage
retraining/0 [default0]:[rank0]:     self._run_sanity_check()
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/trainer.py", line 1052, in _run_sanity_check
retraining/0 [default0]:[rank0]:     val_loop.run()
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/utilities.py", line 178, in _decorator
retraining/0 [default0]:[rank0]:     return loop_run(self, *args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/evaluation_loop.py", line 135, in run
retraining/0 [default0]:[rank0]:     self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/evaluation_loop.py", line 396, in _evaluation_step
retraining/0 [default0]:[rank0]:     output = call._call_strategy_hook(trainer, hook_name, *step_args)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/call.py", line 319, in _call_strategy_hook
retraining/0 [default0]:[rank0]:     output = fn(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 607, in validation_step
retraining/0 [default0]:[rank0]:     out = self.model.validation_step(dataloader_iter, *args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 342, in validation_step
retraining/0 [default0]:[rank0]:     return self._step(
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 429, in _step
retraining/0 [default0]:[rank0]:     return self.forward(
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 279, in forward
retraining/0 [default0]:[rank0]:     microbatch_outputs = step()
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 1149, in __call__
retraining/0 [default0]:[rank0]:     return self.forward_backward_func(
retraining/0 [default0]:[rank0]:   File "/opt/megatron-lm/megatron/core/pipeline_parallel/schedules.py", line 452, in forward_backward_no_pipelining
retraining/0 [default0]:[rank0]:     output_tensor, num_tokens = forward_step(
retraining/0 [default0]:[rank0]:   File "/opt/megatron-lm/megatron/core/pipeline_parallel/schedules.py", line 274, in forward_step
retraining/0 [default0]:[rank0]:     output_tensor, loss_func = forward_step_func(data_iterator, model)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 492, in wrapped_forward_step_func
retraining/0 [default0]:[rank0]:     output_tensor = _forward_step(model, batch)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 760, in wrapped
retraining/0 [default0]:[rank0]:     return attr(*args)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 384, in validation_step
retraining/0 [default0]:[rank0]:     return self.forward_step(batch)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 375, in forward_step
retraining/0 [default0]:[rank0]:     return self.config.forward_step_fn(self, batch)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/ssm.py", line 47, in ssm_forward_step
retraining/0 [default0]:[rank0]:     return model(**forward_args)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
retraining/0 [default0]:[rank0]:     return self._call_impl(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
retraining/0 [default0]:[rank0]:     return forward_call(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 359, in forward
retraining/0 [default0]:[rank0]:     output_tensor = self.module(
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
retraining/0 [default0]:[rank0]:     return self._call_impl(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
retraining/0 [default0]:[rank0]:     return forward_call(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/megatron-lm/megatron/core/distributed/data_parallel_base.py", line 22, in forward
retraining/0 [default0]:[rank0]:     return self.module(*inputs, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
retraining/0 [default0]:[rank0]:     return self._call_impl(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
retraining/0 [default0]:[rank0]:     return forward_call(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/megatron-lm/megatron/core/transformer/module.py", line 178, in forward
retraining/0 [default0]:[rank0]:     outputs = self.module(*inputs, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
retraining/0 [default0]:[rank0]:     return self._call_impl(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
retraining/0 [default0]:[rank0]:     return forward_call(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/megatron-lm/megatron/core/models/mamba/mamba_model.py", line 206, in forward
retraining/0 [default0]:[rank0]:     hidden_states = self.decoder(
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
retraining/0 [default0]:[rank0]:     return self._call_impl(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
retraining/0 [default0]:[rank0]:     return forward_call(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/megatron-lm/megatron/core/ssm/mamba_block.py", line 262, in forward
retraining/0 [default0]:[rank0]:     hidden_states = layer(
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
retraining/0 [default0]:[rank0]:     return self._call_impl(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
retraining/0 [default0]:[rank0]:     return forward_call(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/megatron-lm/megatron/core/ssm/mamba_layer.py", line 105, in forward
retraining/0 [default0]:[rank0]:     mixer_out_with_bias = self.mixer(hidden_states, inference_params=inference_params)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
retraining/0 [default0]:[rank0]:     return self._call_impl(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
retraining/0 [default0]:[rank0]:     return forward_call(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/opt/megatron-lm/megatron/core/ssm/mamba_mixer.py", line 299, in forward
retraining/0 [default0]:[rank0]:     y = mamba_split_conv1d_scan_combined(
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/mamba_ssm/ops/triton/ssd_combined.py", line 930, in mamba_split_conv1d_scan_combined
retraining/0 [default0]:[rank0]:     return MambaSplitConv1dScanCombinedFn.apply(zxbcdt, conv1d_weight, conv1d_bias, dt_bias, A, D, chunk_size, initial_states, seq_idx, dt_limit, return_final_states, activation, rmsnorm_weight, rmsnorm_eps, outproj_weight, outproj_bias, headdim, ngroups, norm_before_gate)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py", line 575, in apply
retraining/0 [default0]:[rank0]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/amp/autocast_mode.py", line 456, in decorate_fwd
retraining/0 [default0]:[rank0]:     return fwd(*args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/mamba_ssm/ops/triton/ssd_combined.py", line 789, in forward
retraining/0 [default0]:[rank0]:     out, out_x, dt_out, dA_cumsum, states, final_states = _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size=chunk_size, D=D, z=z, dt_bias=dt_bias, initial_states=initial_states, seq_idx=seq_idx, dt_softplus=True, dt_limit=dt_limit)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/mamba_ssm/ops/triton/ssd_combined.py", line 312, in _mamba_chunk_scan_combined_fwd
retraining/0 [default0]:[rank0]:     dA_cumsum, dt = _chunk_cumsum_fwd(dt, A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus, dt_limit=dt_limit)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/mamba_ssm/ops/triton/ssd_chunk_state.py", line 675, in _chunk_cumsum_fwd
retraining/0 [default0]:[rank0]:     _chunk_cumsum_fwd_kernel[grid_chunk_cs](
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 345, in <lambda>
retraining/0 [default0]:[rank0]:     return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/runtime/autotuner.py", line 156, in run
retraining/0 [default0]:[rank0]:     timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/runtime/autotuner.py", line 156, in <dictcomp>
retraining/0 [default0]:[rank0]:     timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/runtime/autotuner.py", line 133, in _bench
retraining/0 [default0]:[rank0]:     return do_bench(kernel_call, warmup=self.num_warmups, rep=self.num_reps, quantiles=(0.5, 0.2, 0.8))
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/testing.py", line 106, in do_bench
retraining/0 [default0]:[rank0]:     fn()
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/runtime/autotuner.py", line 114, in kernel_call
retraining/0 [default0]:[rank0]:     self.fn.run(
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 662, in run
retraining/0 [default0]:[rank0]:     kernel = self.compile(
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/compiler/compiler.py", line 282, in compile
retraining/0 [default0]:[rank0]:     next_module = compile_ir(module, metadata)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/compiler.py", line 320, in <lambda>
retraining/0 [default0]:[rank0]:     stages["cubin"] = lambda src, metadata: self.make_cubin(src, metadata, options, self.capability)
retraining/0 [default0]:[rank0]:   File "/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/compiler.py", line 297, in make_cubin
retraining/0 [default0]:[rank0]:     raise RuntimeError(f'Internal Triton PTX codegen error: \n{log}')
retraining/0 [default0]:[rank0]: RuntimeError: Internal Triton PTX codegen error: 
retraining/0 [default0]:[rank0]: ptxas /tmp/tmppa4c09n_.ptx, line 120; error   : Feature '.bf16' requires .target sm_80 or higher
retraining/0 [default0]:[rank0]: ptxas /tmp/tmppa4c09n_.ptx, line 120; error   : Feature 'cvt with .f32.bf16' requires .target sm_80 or higher
retraining/0 [default0]:[rank0]: ptxas fatal   : Ptx assembly aborted due to errors
retraining/0 [default0]:
retraining/0 E0103 05:50:22.565000 17575 torch/distributed/elastic/multiprocessing/api.py:862] failed (exitcode: 1) local_rank: 0 (pid: 17601) of binary: /usr/bin/python
retraining/0 I0103 05:50:22.579000 17575 torch/distributed/elastic/multiprocessing/errors/__init__.py:368] ('local_rank %s FAILED with no error file. Decorate your entrypoint fn with @record for traceback info. See: https://pytorch.org/docs/stable/elastic/errors.html', 0)
retraining/0 Traceback (most recent call last):
retraining/0   File "/usr/local/bin/torchrun", line 33, in <module>
retraining/0     sys.exit(load_entry_point('torch==2.5.0a0+e000cf0ad9.nv24.10', 'console_scripts', 'torchrun')())
retraining/0   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
retraining/0     return f(*args, **kwargs)
retraining/0   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/run.py", line 919, in main
retraining/0     run(args)
retraining/0   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/run.py", line 910, in run
retraining/0     elastic_launch(
retraining/0   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/launcher/api.py", line 138, in __call__
retraining/0     return launch_agent(self._config, self._entrypoint, list(args))
retraining/0   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/launcher/api.py", line 269, in launch_agent
retraining/0     raise ChildFailedError(
retraining/0 torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
retraining/0 ============================================================
retraining/0 nemo_run.core.runners.fdl_runner FAILED
retraining/0 ------------------------------------------------------------
retraining/0 Failures:
retraining/0   <NO_OTHER_FAILURES>
retraining/0 ------------------------------------------------------------
retraining/0 Root Cause (first observed failure):
retraining/0 [0]:
retraining/0   time      : 2025-01-03_05:50:22
retraining/0   host      : bh-Precision-5820-Tower-X-Series
retraining/0   rank      : 0 (local_rank: 0)
retraining/0   exitcode  : 1 (pid: 17601)
retraining/0   error_file: <N/A>
retraining/0   traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
retraining/0 ============================================================
[05:50:23] INFO     Job mamba2_780m_pretraining-ht7tmrvqhgqrgd finished: FAILED                                                            launcher.py:162
                                                                                                                                                          
# The experiment was run with the following tasks: ['mamba2_780m_pretraining']                                                                            
# You can inspect and reconstruct this experiment at a later point in time using:                                                                         
experiment = run.Experiment.from_id("mamba2_780m_pretraining_1735883405")                                                                                 
experiment.status() # Gets the overall status                                                                                                             
experiment.logs("mamba2_780m_pretraining") # Gets the log for the provided task                                                                           
experiment.cancel("mamba2_780m_pretraining") # Cancels the provided task if still running 
``` 

Any tips would be great! Thank you in advance!