fix_untrained_tokens doesn't work with zero-3

[winglian]

Please check that this issue hasn't been reported before.

• I searched previous Bug Reports didn't find any similar reports.

Expected Behavior

n/a

Current behaviour

n/a

Steps to reproduce

n/a

Config yaml

No response

Possible solution

i think this is because the embeddings are sharded

Which Operating Systems are you using?

• Linux
• macOS
• Windows

Python Version

3.11

axolotl branch-commit

main

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this bug has not been reported yet.
• I am using the latest version of axolotl.
• I have provided enough information for the maintainers to reproduce and diagnose the issue.

[williambarberjr]

Ran into this today. For the sake of anyone else arriving from google, I'm using python 3.12 and release version 0.6.0. My config yml has:
fix_untrained_tokens: true

And my error looks like this:

/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/core/trainer_builder.py:444: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `AxolotlTrainer.__init__`. Use `processing_class` instead.
  super().__init__(*_args, **kwargs)
/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/core/trainer_builder.py:444: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `AxolotlTrainer.__init__`. Use `processing_class` instead.
  super().__init__(*_args, **kwargs)
[rank2]: Traceback (most recent call last):
[rank2]:   File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
[rank2]:     return _run_code(code, main_globals, None,
[rank2]:   File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
[rank2]:     exec(code, run_globals)
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/cli/train.py", line 58, in <module>
[rank2]:     fire.Fire(do_cli)
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/fire/core.py", line 135, in Fire
[rank2]:     component_trace = _Fire(component, args, parsed_flag_args, context, name)
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/fire/core.py", line 468, in _Fire
[rank2]:     component, remaining_args = _CallAndUpdateTrace(
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/fire/core.py", line 684, in _CallAndUpdateTrace
[rank2]:     component = fn(*varargs, **kwargs)
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/cli/train.py", line 34, in do_cli
[rank2]:     return do_train(parsed_cfg, parsed_cli_args)
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/cli/train.py", line 47, in do_train
[rank2]:     model, tokenizer = train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/train.py", line 133, in train
[rank2]:     fix_untrained_tokens(model, tokenizer, train_dataset)
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank2]:     return func(*args, **kwargs)
[rank2]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/core/tokenizer_utils.py", line 51, in fix_untrained_tokens
[rank2]:     min_size = min(embedding_matrix.shape[1], lm_head_matrix.shape[1])
[rank2]: IndexError: tuple index out of range
[rank5]: Traceback (most recent call last):
[rank5]:   File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
[rank5]:     return _run_code(code, main_globals, None,
[rank5]:   File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
[rank5]:     exec(code, run_globals)
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/cli/train.py", line 58, in <module>
[rank5]:     fire.Fire(do_cli)
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/fire/core.py", line 135, in Fire
[rank5]:     component_trace = _Fire(component, args, parsed_flag_args, context, name)
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/fire/core.py", line 468, in _Fire
[rank5]:     component, remaining_args = _CallAndUpdateTrace(
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/fire/core.py", line 684, in _CallAndUpdateTrace
[rank5]:     component = fn(*varargs, **kwargs)
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/cli/train.py", line 34, in do_cli
[rank5]:     return do_train(parsed_cfg, parsed_cli_args)
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/cli/train.py", line 47, in do_train
[rank5]:     model, tokenizer = train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/train.py", line 133, in train
[rank5]:     fix_untrained_tokens(model, tokenizer, train_dataset)
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank5]:     return func(*args, **kwargs)
[rank5]:   File "/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/core/tokenizer_utils.py", line 51, in fix_untrained_tokens
[rank5]:     min_size = min(embedding_matrix.shape[1], lm_head_matrix.shape[1])
[rank5]: IndexError: tuple index out of range
/ephemeral/axolotl_env/lib/python3.10/site-packages/axolotl/core/trainer_builder.py:444: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `AxolotlTrainer.__init__`. Use `processing_class` instead.
  super().__init__(*_args, **kwargs)

Didn't try it but I assume using zero2.json would resolve this?