Pull in main from argonne-lcf/Megatron-DeepSpeed

[saforem2]

Summary by Sourcery

Refactor the training script for improved logging and modularity, introduce new features such as support for additional optimizers and tokenizers, and enhance the data processing pipeline. Add new scripts for setting up and running training jobs on ALCF systems, and update documentation to reflect these changes.

New Features:

• Introduce a new logging mechanism using the log object from the logging module, replacing the previous print_rank_0 calls for improved logging consistency.
• Add support for multiple optimizers including galore_adamw, adafactor, and adam8bit, providing more flexibility in training configurations.
• Implement a new tokenizer type Llama2Tokenizer for handling specific tokenization needs.
• Add a new script pretrain_llama.py for pretraining Llama models, which includes setup for distributed training and logging with WandB.
• Introduce a new helper script ALCF/helpers.sh for setting up and launching training jobs on ALCF systems, including environment setup and job configuration.

Bug Fixes:

• Fix issues with the get_ltor_masks_and_position_ids function to ensure correct mask and position ID generation.
• Resolve compatibility issues with the torch and deepspeed versions by updating the setup and installation scripts.

Enhancements:

• Refactor the training script to improve modularity and logging, including the use of decorators for logging and timing functions.
• Enhance the data loading process by introducing a BuildConcatDataset class for efficient dataset handling and concatenation.
• Improve the optimizer setup by allowing more granular control over parameter groups and conditions for weight decay and learning rate scaling.
• Optimize the data processing pipeline by adding support for sequence parallelism and curriculum learning adjustments.

Build:

• Add a new build script make_data to ensure the megatron/data/helpers.cpp is compiled before training.

Deployment:

• Add a new script ALCF/README.md to guide users on setting up and running training jobs on ALCF systems, including detailed instructions for different environments.

Documentation:

• Update documentation to include new features and enhancements, particularly around the new tokenizer and optimizer options.

Chores:

• Organize and clean up the codebase by removing unused imports and redundant code blocks.

[saforem2]

LGTM 👍