No longer ignore --iteration when passed to train.py

[haileyschoelkopf]

We currently resume training using the checkpoint tag provided in the "latest" file within the save folder, following Deepspeed's default behavior. However, we want to be able to override this "latest" tag when launching a job when, for example, resuming several checkpoints back from a corrupted run.

An --iteration flag that can be passed to deepy.py already exists, and can be successfully used to evaluate a model from any step with evaluate.py , ignoring the "latest" file. However, though this flag properly sets neox_args.iteration at initialization, when initializing a model we ignore this because it is not passed to setup_model_and_optimizer.

This PR resolves this, allowing users to successfully use --iteration alongside train.py.

@Quentin-Anthony @ShivanshuPurohit

[StellaAthena]

Tested and found to work, but didn’t stress test unusual cases or potential human errors. Approving, but I’ll leave it to @Quentin-Anthony to decide if it should be merged or tested further.

Also, I noticed we are passing use_cache=False… do we never want to use cache?

[haileyschoelkopf]

re: use_cache=False: yes, in training we don't ever want or need cache. If use_cache is ever true it gets set to false when enabling train mode:

gpt-neox/megatron/model/gpt2_model.py

Line 331 in 038b011

recursive_setattr(self.forward_funcs, "use_cache", False)

[Quentin-Anthony]

I'm happy with this. Merging.