NVIDIA-NeMo/Automodel

Is your feature request related to a problem? Please describe. See here see https://github.com/marcelroed/gigatoken#gigatoken for context ; Describe the solution you'd like A clear and concise description of what you want to happen. Describe alternatives you've considered A clear and concise description of any alternative solutions or features you've considered. Additional context Add any other context or screenshots about the feature request here.

Context Follow-up to AM-730. The successful scoped release CI run for nemotron_nano_v3_singlegpu_lora proved the AM-730 activation-checkpointing fix, but loading the 58.82 GiB base checkpoint took 570.10 seconds: Disk read: 5.78 seconds HF-to-AutoModel state conversion and model copy (currently logged together as distribute): 564.32 seconds Overall effective rate: 0.10 GiB/s The job used an exclusive EOS node with 224 allocated CPUs, while Slurm reported CPUs/Task=1. The run did not log process affinity or PyTorch intra-op/inter-op thread counts, so a one-CPU restriction is a hypothesis, not an established root cause. The current path CPU-merges thousands of split expert tensors into grouped MoE weights, performs repeated garbage collection/cache cleanup, and then copies the full state into the single-GPU model. The JoC harness supports CPUS_PER_TASK, but AutoModel's recipe CI generator does not currently expose it. Investigation Log os.sched_getaffinity(0), torch.get_num_threads(),

https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro

Is your feature request related to a problem? Please describe. AM uses .map(fn) to transform datasets (e.g., apply template, tokenize etc). Unfortunately, this happens during startup and runs over the whole dataset. In addition, the current caching functionality often fails, incurring additional delay on repeated runs. 2026-04-20 12:09:22 | INFO | httpx | HTTP Request: HEAD https://huggingface.co/akoumpa/Devstral-Small-2-24B-Instruct-2512-BF16/resolve/main/config.json "HTTP/1.1 307 Temporary Redirect" 2026-04-20 12:09:22 | INFO | httpx | HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/akoumpa/Devstral-Small-2-24B-Instruct-2512-BF16/6d6f531753869fcd8cbbaefa7ee94b93966d2d04/config.json "HTTP/ 1.1 200 OK" 2026-04-20 12:09:23 | INFO | nemo_automodel._transformers.auto_tokenizer | Using custom tokenizer MistralCommonBackend for model type 'mistral3' 2026-04-20 12:09:23 | INFO | httpx | HTTP Request: GET https://huggingface.co/api/models/akoumpa/Devstral-Small-2-24B-Instru

Describe the bug In particular for #1631, after fixing the tied weights issue, I run into another issue where the default collator was creating a "padding_mask" although the model is unable to consume it, thus crashing the run. 1631 included a filtering step to remove the padding_mask, but it would be better not to have to do any filtering, as it pollutes/complicates the code. Steps/Code to reproduce bug Run the yaml that I have attached to 1631 Expected behavior Run without crashing / having to filter the batch. Additional context Add any other context about the problem here.