Generic tracker API and implementation of Aimstack tracker

[dushyantbehl]

Description of the change

1. This PR adds a tracker API with Aimstack as the default tracker. This is simple plug and play architecture which can support multiple trackers.
2. The tracker config is now taken as command line arguments (making it easier for any automation to pass tracker arguments).
3. With the new API I have added support to track any additional metrics of interest.
4. As an example I have added one single line to track model_load_time in AIM possibly fixing Add support for collecting metrics programmatically #33

Related issue number

1. Add support for registering custom AIM metadata #34
2. Add support for collecting metrics programmatically #33

How to verify the PR

Example of how new api can be invoked

 torchrun --nnodes=1 --nproc_per_node=2 --master_port=1234 tuning/sft_trainer.py --tokenizer_name_or_path ${MODEL_PATH} --model_name_or_path ${MODEL_PATH} --data_path ${DATA_PATH} --use_peft --bf16 True --output_dir ${OUTPUT_PATH} --num_train_epochs 1 --per_device_train_batch_size 1 --per_device_eval_batch_size 1 --gradient_accumulation_steps 8 --evaluation_strategy "no" --save_strategy "steps" --save_steps 2000 --save_total_limit 1 --learning_rate 2e-5 --weight_decay 0. --warmup_ratio 0.03 --lr_scheduler_type "cosine" --logging_steps 1 --fsdp "full_shard auto_wrap" --fsdp_config tuning/config/fsdp_config.json --bf16 True --response_template "\n### Response:" --dataset_text_field "output" --tracker aim --aim_repo /data/aim --experiment sft-llama7b-test 

Was the PR tested

• I have added >=1 unit test(s) for every new method I have added.
• I have ensured all unit tests pass

[dushyantbehl]

@Ssukriti @alex-jw-brooks some more changes and fixes as requested by @kmehant

PR is ready to be reviewed. @tharapalanivel I will run some more e2e tests and bump you, if you have any comments in the meantime on the design, happy to take them.

[tharapalanivel]

@dushyantbehl this is great, thank you! One thing we should be mindful of is that the image build process still works and picks up the new changes. For this we should either change build/launch_training.py to use sft_trainer.main() or move the tracker/callback logic to build/utils.process_launch_training_args() so launch_training.py will use it automatically. Images can be built off of the branch and tested locally to check that everything works as expected, you can find docs for this here. I know this might involve mocking/patching a few things but can we look into whether we can add any meaningful tests for this please? Let me know if you have any questions, thanks!

[dushyantbehl]

Thanks @tharapalanivel should be easy to change it like that.

[Ssukriti]

comments only on default value of arguments.

By default I want to ensure existing behavior is retained - any users that call .train() or main() function using command line, should see log files like before. So filelogger has to be default for train() and through command line.
Only if aim is explictly passed, it should be used.