Adding tabs for different set of Accelerate's features and content for large scale training features (#2)

- Upload 18 files (e5cadf9d99a28f3f5354729e0d2c8d93a90bd7e0)


Co-authored-by: Sourab Mangrulkar <smangrul@users.noreply.huggingface.co>

code_samples/base/accelerate

@@ -0,0 +1,17 @@
+<pre>
+from accelerate import Accelerator
+accelerator = Accelerator()
+train_dataloader, model, optimizer scheduler = accelerator.prepare(
+        dataloader, model, optimizer, scheduler
+)
+
+model.train()
+for batch in train_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+</pre>

code_samples/base/basic

@@ -0,0 +1,31 @@
+##
+<pre>
++from accelerate import Accelerator
++accelerator = Accelerator()
++dataloader, model, optimizer scheduler = accelerator.prepare(
++        dataloader, model, optimizer, scheduler
++)
+
+for batch in dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+-    loss.backward()
++    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()</pre>
+##
+Everything around `accelerate` occurs with the `Accelerator` class. To use it, first make an object.
+Then call `.prepare` passing in the PyTorch objects that you would normally train with. This will
+return the same objects, but they will be on the correct device and distributed if needed. Then
+you can train as normal, but instead of calling `loss.backward()` you call `accelerator.backward(loss)`.
+Also note that you don't need to call `model.to(device)` or `inputs.to(device)` anymore, as this
+is done automatically by `accelerator.prepare()`.
+
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/basic_tutorials/migration" target="_blank">Migrating to 🤗 Accelerate</a>
+- <a href="https://huggingface.co/docs/accelerate/package_reference/accelerator" target="_blank">The Accelerator</a>

code_samples/base/calculating_metrics

@@ -0,0 +1,51 @@
+##
+<pre>
+import evaluate
++from accelerate import Accelerator
++accelerator = Accelerator()
++train_dataloader, eval_dataloader, model, optimizer, scheduler = (
++    accelerator.prepare(
++        train_dataloader, eval_dataloader, 
++        model, optimizer, scheduler
++    )
++)
+metric = evaluate.load("accuracy")
+for batch in train_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+
+model.eval()
+for batch in eval_dataloader:
+    inputs, targets = batch
+-    inputs = inputs.to(device)
+-    targets = targets.to(device)
+    with torch.no_grad():
+        outputs = model(inputs)
+    predictions = outputs.argmax(dim=-1)
++    predictions, references = accelerator.gather_for_metrics(
++        (predictions, references)
++    )
+    metric.add_batch(
+        predictions = predictions,
+        references = references
+    )
+print(metric.compute())</pre>
+
+##
+When calculating metrics on a validation set, you can use the `Accelerator.gather_for_metrics`
+method to gather the predictions and references from all devices and then calculate the metric on the gathered values. 
+This will also *automatically* drop the padded values from the gathered tensors that were added to ensure 
+that all tensors have the same length. This ensures that the metric is calculated on the correct values.
+##
+To learn more checkout the related documentation:
+
+- <a href="https://huggingface.co/docs/accelerate/en/quicktour#distributed-evaluation" target="_blank">Quicktour - Calculating metrics</a>
+- <a href="https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.gather_for_metrics" target="_blank">API reference</a>
+- <a href="https://github.com/huggingface/accelerate/blob/main/examples/by_feature/multi_process_metrics.py" target="_blank">Example script</a>

code_samples/base/checkpointing

@@ -0,0 +1,29 @@
+##
+<pre>
+from accelerate import Accelerator
+accelerator = Accelerator()
+dataloader, model, optimizer scheduler = accelerator.prepare(
+        dataloader, model, optimizer, scheduler
+)
+
+for batch in dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
++accelerator.save_state("checkpoint_dir")
++accelerator.load_state("checkpoint_dir")</pre>
+##
+To save or load a checkpoint in, `Accelerator` provides the `save_state` and `load_state` methods.
+These methods will save or load the state of the model, optimizer, scheduler, as well as random states and
+any custom registered objects from the main process on each device to a passed in folder. 
+**This API is designed to save and resume training states only from within the same python script or training setup.**
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/usage_guides/checkpoint" target="_blank">Saving and loading training states</a>
+- <a href="https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state" target="_blank">`save_state` API reference</a>
+- <a href="https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.load_state" target="_blank">`load_state` API reference</a>
+- <a href="https://github.com/huggingface/accelerate/blob/main/examples/by_feature/checkpointing.py" target="_blank">Example script</a>

code_samples/base/experiment_tracking

@@ -0,0 +1,32 @@
+##
+<pre>
+from accelerate import Accelerator
+-accelerator = Accelerator()
++accelerator = Accelerator(log_with="wandb")
+train_dataloader, model, optimizer scheduler = accelerator.prepare(
+        dataloader, model, optimizer, scheduler
+)
++accelerator.init_trackers()
+model.train()
+for batch in train_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
++    accelerator.log({"loss":loss})
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
++accelerator.end_training()
+</pre>
+##
+To use experiment trackers with `accelerate`, simply pass the desired tracker to the `log_with` parameter
+when building the `Accelerator` object. Then initialize the tracker(s) by running `Accelerator.init_trackers()`
+passing in any configurations they may need. Afterwards call `Accelerator.log` to log a particular value to your tracker.
+At the end of training call `accelerator.end_training()` to call any finalization functions a tracking library
+may need automatically.
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/usage_guides/tracking" target="_blank">Using experiment trackers</a>
+- <a href="https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.log" target="_blank">Accelerator API Reference</a>
+- <a href="https://huggingface.co/docs/accelerate/package_reference/tracking" target="_blank">Tracking API Reference</a>

code_samples/base/gradient_accumulation

@@ -0,0 +1,33 @@
+##
+<pre>
+from accelerate import Accelerator
+accelerator = Accelerator(
++    gradient_accumulation_steps=2,
+)
+dataloader, model, optimizer scheduler = accelerator.prepare(
+        dataloader, model, optimizer, scheduler
+)
+
+for batch in dataloader:
++  with accelerator.accumulate(model):
+      optimizer.zero_grad()
+      inputs, targets = batch
+      outputs = model(inputs)
+      loss = loss_function(outputs, targets)
+      accelerator.backward(loss)
+      optimizer.step()
+      scheduler.step()</pre>
+
+##
+When performing gradient accumulation in a distributed setup, there are many opportunities for efficiency mistakes
+to occur. `Accelerator` provides a context manager that will take care of the details for you and ensure that the
+model is training correctly. Simply wrap the training loop in the `Accelerator.accumulate` context manager
+while passing in the model you are training on and during training the gradients will accumulate and synchronize
+automatically when needed.
+
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/usage_guides/gradient_accumulation" target="_blank">Performing gradient accumulation</a>
+- <a href="https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.accumulate" target="_blank">API reference</a>
+- <a href="https://github.com/huggingface/accelerate/blob/main/examples/by_feature/gradient_accumulation.py" target="_blank">Example script</a>
+- <a href="https://github.com/huggingface/accelerate/blob/main/examples/by_feature/automatic_gradient_accumulation.py" target="_blank">Performing automatic gradient accumulation example script</a>

code_samples/base/initial

@@ -0,0 +1,11 @@
+<pre>
+for batch in dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()</pre>

code_samples/base/initial_with_metrics

@@ -0,0 +1,27 @@
+<pre>
+import evaluate
+metric = evaluate.load("accuracy")
+for batch in train_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+
+model.eval()
+for batch in eval_dataloader:
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    with torch.no_grad():
+        outputs = model(inputs)
+    predictions = outputs.argmax(dim=-1)
+    metric.add_batch(
+        predictions = predictions,
+        references = references
+    )
+print(metric.compute())</pre>

code_samples/large_scale_training/aws_sagemaker

@@ -0,0 +1,77 @@
+##
+Run  `accelerate config` on and answer the questionnaire accordingly. 
+Below is an example yaml for running code remotely on AWS SageMaker. Replace placeholder `xxxxx` with 
+appropriate values.
+
+<pre>
+base_job_name: accelerate-sagemaker-1
+compute_environment: AMAZON_SAGEMAKER
+distributed_type: 'NO'
+dynamo_backend: 'NO'
+ec2_instance_type: ml.p3.2xlarge
+gpu_ids: all
+iam_role_name: xxxxx
+mixed_precision: 'no'
+num_machines: 1
+profile: xxxxx
+py_version: py38
+pytorch_version: 1.10.2
+region: us-east-1
+transformers_version: 4.17.0
+use_cpu: false
+</pre>
+##
+<pre>
+from accelerate import Accelerator
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="sample task")
+
+    parser.add_argument(
+        "--pad_to_max_length",
+-        action="store_true",
++        type=bool,
++        default=False,
+        help="If passed, pad all samples to `max_length`. Otherwise, dynamic padding is used.",
+    )
+
+    ...
+
+  
++ def main():
+      accelerator = Accelerator()
+
+      model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+          model, optimizer, training_dataloader, scheduler
+      )
+
+      for batch in training_dataloader:
+          optimizer.zero_grad()
+          inputs, targets = batch
+          outputs = model(inputs)
+          loss = loss_function(outputs, targets)
+          accelerator.backward(loss)
+          optimizer.step()
+          scheduler.step()
+
+-    torch.save('/opt/ml/model`)
++    accelerator.save('/opt/ml/model')
+
++ if __name__ == "__main__":
++     main()
+</pre>
+Launching a script using default accelerate config file looks like the following:
+```
+accelerate launch {script_name.py} {--arg1} {--arg2} ...
+```
+##
+SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to specify type as bool in your script and provide an explicit True or False value for this hyperparameter. An example for the same is shown above for `pad_to_max_length` argument. Another important point is to save all the output artifacts to `/opt/ml/model` or use `os.environ["SM_MODEL_DIR"]` as your save directory. After training, artifacts in this directory are uploaded to S3, an example is shown in above code snippet.
+
+You can provide custom docker image, input channels pointing to S3 data locations and use SageMaker metrics logging
+as part of advanced features. Please refer <a href="https://github.com/huggingface/notebooks/tree/main/sagemaker/22_accelerate_sagemaker_examples" target="_blank">Examples showcasing AWS SageMaker integration of 🤗 Accelerate</a>
+
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/usage_guides/sagemaker" target="_blank">How to use 🤗 Accelerate with SageMaker</a>
+- <a href="https://github.com/huggingface/notebooks/tree/main/sagemaker/22_accelerate_sagemaker_examples" target="_blank">Examples showcasing AWS SageMaker integration of 🤗 Accelerate</a>
+- <a href="https://huggingface.co/docs/accelerate/main/en/package_reference/cli" target="_blank">The Accelerate CLI</a>

code_samples/large_scale_training/deepspeed

@@ -0,0 +1,101 @@
+##
+Run  `accelerate config` and answer the questionnaire accordingly. 
+Below is an example yaml for mixed-precision training using DeepSpeed ZeRO Stage-3 with CPU offloading on 8 GPUs.
+<pre>
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: cpu
+  offload_param_device: cpu
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+machine_rank: 0
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: fp16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+use_cpu: false
+</pre>
+##
+<pre>
+  from accelerate import Accelerator
+  
++ def main():
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
+    ...
+
+    generated_tokens = accelerator.unwrap_model(model).generate(
+                    batch["input_ids"],
+                    attention_mask=batch["attention_mask"],
+                    **gen_kwargs,
++                    synced_gpus=True #required for ZeRO Stage 3
+                )
+    ...
+
+    accelerator.unwrap_model(model).save_pretrained(
+            args.output_dir,
+            is_main_process=accelerator.is_main_process,
+            save_function=accelerator.save,
++            state_dict=accelerator.get_state_dict(model), #required for ZeRO Stage 3
+        )
+
+    ...
+
++ if __name__ == "__main__":
++     main()
+</pre>
+
+Launching a script using default accelerate config file looks like the following:
+```
+accelerate launch {script_name.py} {--arg1} {--arg2} ...
+```
+
+Alternatively, you can use `accelerate launch` with right config params for multi-gpu training as shown below
+```
+accelerate launch \
+    --use_deepspeed \
+    --num_processes=8 \
+    --mixed_precision=fp16 \
+    --zero_stage=3 \
+    --gradient_accumulation_steps=1 \
+    --gradient_clipping=1 \
+    --zero3_init_flag=True \
+    --zero3_save_16bit_model=True \
+    --offload_optimizer_device=cpu \
+    --offload_param_device=cpu \
+    {script_name.py} {--arg1} {--arg2} ...
+```
+
+##
+For core DeepSpeed features supported via accelerate config file, no changes are required for ZeRO Stages 1 and 2. For ZeRO Stage-3, transformers' `generate` function requires `synced_gpus=True` and `save_pretrained` requires the `state_dict` param due to the fact that model parameters are sharded across the GPUs.
+
+For advanced users who like granular control via DeepSpeed config file, it is supported wherein you can pass its loaction when running `accelerate config` command. You can also specify values of most of the fields in DeepSpeed config file as `auto` and they are filled automatically via the arguments of `accelerate launch` command and `accelerator.prepare` call thereby making life simple for users. Please refer docs on <a href="https://huggingface.co/docs/accelerate/usage_guides/deepspeed#deepspeed-config-file" target="_blank">DeepSpeed Config File</a>
+
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/usage_guides/deepspeed" target="_blank">How to use DeepSpeed</a>
+- <a href="https://huggingface.co/blog/accelerate-deepspeed" target="_blank">Accelerate Large Model Training using DeepSpeed</a>
+- <a href="https://huggingface.co/docs/accelerate/package_reference/deepspeed" target="_blank">DeepSpeed Utilities</a>

code_samples/large_scale_training/megatron-lm

@@ -0,0 +1,119 @@
+##
+Run  `accelerate config` and answer the questionnaire accordingly. 
+Below is an example yaml for BF16 mixed-precision training using Megatron-LM with DPxTPxPP=2x2x2 degrees on 8 GPUs. (DP-Data Parallelism, PP-Pipeline Parallelism, TP-Tensor Parallelism). It is also using Sequence Parallelism and selective activation checkpointing along with sharded optimizer.
+<pre>
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MEGATRON_LM
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+machine_rank: 0
+main_training_function: main
+megatron_lm_config:
+  megatron_lm_gradient_clipping: 1.0
+  megatron_lm_num_micro_batches: 2
+  megatron_lm_pp_degree: 2
+  megatron_lm_recompute_activations: true
+  megatron_lm_sequence_parallelism: true
+  megatron_lm_tp_degree: 2
+  megatron_lm_use_distributed_optimizer: true
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+use_cpu: false
+</pre>
+##
+<pre>
+  from accelerate import Accelerator
+  
++ def main():
+    accelerator = Accelerator()
+
+    ...
+
+-    lr_scheduler = get_scheduler(
+-        name=args.lr_scheduler_type,
++    lr_scheduler = accelerate.utils.MegatronLMDummyScheduler(
+        optimizer=optimizer,
+        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+    )
+
+
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    total_batch_size = (
+-            args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
++            accelerator.state.megatron_lm_plugin.global_batch_size
+        )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
+    ...
+
+    # in eval loop
+    for step, batch in enumerate(eval_dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+        loss = outputs.loss
+-        losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))
++        losses.append(loss) # For Megatron-LM, the losses are already averaged across the data parallel group
+-    losses = torch.cat(losses)
++    losses = torch.tensor(losses)
+    eval_loss = torch.mean(losses)
+    perplexity = math.exp(eval_loss)
+    logger.info(f"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}")
+
++    accelerator.save_state(output_dir)
+
++ if __name__ == "__main__":
++     main()
+</pre>
+
+Launching a script using default accelerate config file looks like the following:
+```
+accelerate launch {script_name.py} {--arg1} {--arg2} ...
+```
+
+Alternatively, you can use `accelerate launch` with right config params for multi-gpu training as shown below
+```
+accelerate launch \
+    --use_megatron_lm \
+    --num_processes=8 \
+    --mixed_precision=bf16 \
+    --megatron_lm_tp_degree=2 \
+    --megatron_lm_pp_degree=2 \
+    --megatron_lm_num_micro_batches=2 \
+    --megatron_lm_sequence_parallelism=true \
+    --megatron_lm_recompute_activations=true \
+    --megatron_lm_use_distributed_optimizer=true \
+    {script_name.py} {--arg1} {--arg2} ...
+```
+
+##
+For Megatron-LM, the supported models Transformers GPT2, Megatron-BERT and T5 models covering Decoder only, Encode only and Encoder-Decoder model classes. Given the complexity of the features of Megatron-LM, 4 changes that are required to get started are:
+    1. Using `accelerate.utils.MegatronLMDummyScheduler` as Megatron-LM uses its own implementation of Optimizer, the corresponding scheduler compatible with it needs to be used.
+    2. Getting the details of the total batch size now needs to be cognization of tensor and pipeline parallel sizes.
+    3. Losses are already averaged across the data parallel group
+    4. save the model using `accelerator.save_state` instead of transformers `from_pretrianed`
+
+These changes have been highlited in the code snippet above.
+
+Megatron-LM intergration supports many advanced features such as ability to leverage custom train step, using Megatron-LM indexed datasets, checkpoint reshaping and interoperabiloity utilities, `megatron_generate` function for text generation using Tensor and Pipeline Parallelism and support for ROPE/ALibi Positional embeddings and Multi-Query Attention. However, these require more changes owing to the complexity; worth it for getting the highest performance.
+
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/usage_guides/megatron_lm" target="_blank">How to use Megatron-LM</a>
+- <a href="https://github.com/pacman100/accelerate-megatron-test" target="_blank">Examples showcasing the Megatron-LM integration of Accelerate</a>

code_samples/large_scale_training/multi_gpu

@@ -0,0 +1,60 @@
+##
+Run  `accelerate config` and answer the questionnaire accordingly. 
+Below is an example yaml for using multi-gpu training with 4 GPUs.
+<pre>
+compute_environment: LOCAL_MACHINE                                                                                             
+deepspeed_config: {}
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+use_cpu: false</pre>
+##
+<pre>
+  from accelerate import Accelerator
+  
++ def main():
+    accelerator = Accelerator()
+
+    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
++ if __name__ == "__main__":
++     main()
+</pre>
+
+Launching a script using default accelerate config file looks like the following:
+```
+accelerate launch {script_name.py} {--arg1} {--arg2} ...
+```
+
+Alternatively, you can use `accelerate launch` with right config params for multi-gpu training as shown below
+```
+accelerate launch --multi_gpu --num_processes=4 {script_name.py} {--arg1} {--arg2} ...
+```
+
+##
+Using this feature involves no changes to the code apart from the ones mentioned in the tab `Simplify your code and improve efficieny`.
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/main/en/basic_tutorials/launch" target="_blank">Launching distributed code</a>
+- <a href="https://huggingface.co/docs/accelerate/main/en/package_reference/cli" target="_blank">The Accelerate CLI</a>

code_samples/large_scale_training/multi_node_multi_gpu

@@ -0,0 +1,89 @@
+##
+Run  `accelerate config` on and answer the questionnaire accordingly. 
+Below is an example yaml for using multi-gpu training with 4 GPUs on 2 nodes/machines.
+
+On Node/Machine 1:
+<pre>
+compute_environment: LOCAL_MACHINE                                                                                             
+deepspeed_config: {}
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: all
+machine_rank: 0
+main_process_ip: 192.168.20.1
+main_process_port: 8080
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: 'no'
+num_machines: 2
+num_processes: 8
+rdzv_backend: static
+same_network: true
+use_cpu: false
+</pre>
+
+On Node/Machine 2:
+<pre>
+compute_environment: LOCAL_MACHINE                                                                                             
+deepspeed_config: {}
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+gpu_ids: all
+-machine_rank: 0
++machine_rank: 1
+main_process_ip: 192.168.20.1
+main_process_port: 8080
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: 'no'
+num_machines: 2
+num_processes: 8
+rdzv_backend: static
+same_network: true
+use_cpu: false
+</pre>
+##
+<pre>
+  from accelerate import Accelerator
+  
++ def main():
+      accelerator = Accelerator()
+
+      model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+          model, optimizer, training_dataloader, scheduler
+      )
+
+      for batch in training_dataloader:
+          optimizer.zero_grad()
+          inputs, targets = batch
+          outputs = model(inputs)
+          loss = loss_function(outputs, targets)
+          accelerator.backward(loss)
+          optimizer.step()
+          scheduler.step()
+
++ if __name__ == "__main__":
++     main()
+</pre>
+
+Launching a script using default accelerate config file looks like the following:
+```
+accelerate launch {script_name.py} {--arg1} {--arg2} ...
+```
+
+Alternatively, you can use `accelerate launch` with right config params for multi-gpu training as shown below. Replace `{node_number}` with appropriate number.
+```
+accelerate launch --multi_gpu --num_machines=2 --num_processes=8 --main_process_ip="192.168.20.1" --main_process_port=8080
+ --machine_rank={node_number} {script_name.py} {--arg1} {--arg2} ...
+```
+
+##
+Using this feature involves no changes to the code apart from the ones mentioned in the tab `Simplify your code and improve efficieny`.
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/main/en/basic_tutorials/launch" target="_blank">Launching distributed code</a>
+- <a href="https://huggingface.co/docs/accelerate/main/en/package_reference/cli" target="_blank">The Accelerate CLI</a>

code_samples/large_scale_training/pytorch_fsdp

@@ -0,0 +1,80 @@
+##
+Run  `accelerate config` and answer the questionnaire accordingly. 
+Below is an example yaml for BF16 mixed-precision training using PyTorch FSDP with CPU offloading on 8 GPUs.
+<pre>
+compute_environment: LOCAL_MACHINE                                                                                             
+deepspeed_config: {}
+distributed_type: FSDP
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: 1
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: T5Block
+machine_rank: 0
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+use_cpu: false
+</pre>
+##
+<pre>
+  from accelerate import Accelerator
+  
++ def main():
+    accelerator = Accelerator()
+
+    model = accelerator.prepare(model)
+
+    optimizer, training_dataloader, scheduler = accelerator.prepare(
+        optimizer, training_dataloader, scheduler
+    )
+
+    for batch in training_dataloader:
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+
+    ...
+
++ if __name__ == "__main__":
++     main()
+</pre>
+
+Launching a script using default accelerate config file looks like the following:
+```
+accelerate launch {script_name.py} {--arg1} {--arg2} ...
+```
+
+Alternatively, you can use `accelerate launch` with right config params for multi-gpu training as shown below
+```
+accelerate launch \
+    --use_fsdp \
+    --num_processes=8 \
+    --mixed_precision=bf16 \
+    --fsdp_sharding_strategy=1 \
+    --fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=T5Block \
+    --fsdp_offload_params=true \
+    {script_name.py} {--arg1} {--arg2} ...
+```
+
+##
+For PyTorch FDSP, you need to prepare the model first before preparing the optimizer since FSDP will shard parameters in-place and this will break any previously initialized optimizers. Same in outlined in the above code snippet. For transformer models, please use `TRANSFORMER_BASED_WRAP` auto wrap policy as shown in the config above. 
+
+
+##
+To learn more checkout the related documentation:
+- <a href="https://huggingface.co/docs/accelerate/usage_guides/fsdp" target="_blank">How to use FSDP</a>
+- <a href="https://huggingface.co/blog/pytorch-fsdp" target="_blank">Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel</a>

src/app.py

@@ -4,47 +4,140 @@ from template import get_templates
 
 templates = get_templates()
 
-def change(inp):
+
+def change(inp, textbox):
     """Based on an `inp`, render and highlight the appropriate code sample.
 
     Args:
         inp (`str`):
             The input button from the interface.
+        textbox (`str`):
+            The textbox specifying the tab name from the interface.
 
     Returns:
         `tuple`: A tuple of the highlighted code diff, and the title for the section.
     """
-    code, explanation, docs = get_text(inp)
-    if inp == "Basic":
-        return (highlight(code), "## Accelerate Code (Base Integration)", explanation, docs)
-    elif inp == "Calculating Metrics":
-        return (highlight(code), f"## Accelerate Code ({inp})", explanation, docs)
-    else:
-        return (highlight(code), f"## Accelerate Code ({inp})", explanation, docs)
+    if textbox == "base":
+        code, explanation, docs = get_text(inp, textbox)
+        if inp == "Basic":
+            return (highlight(code), "## Accelerate Code (Base Integration)", explanation, docs)
+        elif inp == "Calculating Metrics":
+            return (highlight(code), f"## Accelerate Code ({inp})", explanation, docs)
+        else:
+            return (highlight(code), f"## Accelerate Code ({inp})", explanation, docs)
+    elif textbox == "large_scale_training":
+        config, code, explanation, docs = get_text(inp, textbox)
+        return (highlight(config), highlight(code), f"## Accelerate Code ({inp})", explanation, docs)
 
-default = change("Basic")
 
-with gr.Blocks() as demo:
+default = change("Basic", "base")
+
+
+def base_features(textbox):
+    # textbox.value = "base"
     inp = gr.Radio(
-        ["Basic", "Calculating Metrics", "Checkpointing", "Experiment Tracking", "Gradient Accumulation"], 
+        ["Basic", "Calculating Metrics", "Checkpointing", "Experiment Tracking", "Gradient Accumulation"],
         label="Select a feature you would like to integrate",
-        value="Basic"
+        value="Basic",
     )
     with gr.Row():
         with gr.Column():
             feature = gr.Markdown("## Accelerate Code")
             out = gr.Markdown(default[0])
     with gr.Row():
-        with gr.Column():    
+        with gr.Column():
             gr.Markdown("## Explanation")
             explanation = gr.Markdown(default[2])
     with gr.Row():
-        with gr.Column():    
+        with gr.Column():
             gr.Markdown("## Documentation Links")
             docs = gr.Markdown(default[3])
-    inp.change(
-        fn=change, 
-        inputs=inp,
-        outputs=[out, feature, explanation, docs]
+    inp.change(fn=change, inputs=[inp, textbox], outputs=[out, feature, explanation, docs])
+
+
+def large_scale_training(textbox):
+    # textbox.value = "large_scale_training"
+    inp = gr.Radio(
+        ["Multi GPU", "Multi Node Multi GPU", "AWS SageMaker", "DeepSpeed", "PyTorch FSDP", "Megatron-LM"],
+        label="Select a feature you would like to integrate",
+        value="Basic",
     )
-demo.launch()
+    with gr.Row():
+        with gr.Column():
+            feature = gr.Markdown("## Accelerate Config")
+            config = gr.Markdown("")
+    with gr.Row():
+        with gr.Column():
+            feature = gr.Markdown("## Accelerate Code")
+            out = gr.Markdown("")
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## Explanation")
+            explanation = gr.Markdown("")
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## Documentation Links")
+            docs = gr.Markdown("")
+    inp.change(fn=change, inputs=[inp, textbox], outputs=[config, out, feature, explanation, docs])
+
+
+# def big_model_inference():
+#     inp = gr.Radio(
+#         ["Accelerate's Big Model Inference",], # "DeepSpeed ZeRO Stage-3 Offload"
+#         label="Select a feature you would like to integrate",
+#         value="Basic",
+#     )
+#     with gr.Row():
+#         with gr.Column():
+#             feature = gr.Markdown("## Accelerate Code")
+#             out = gr.Markdown(default[0])
+#     with gr.Row():
+#         with gr.Column():
+#             gr.Markdown(default[1])
+#             explanation = gr.Markdown(default[2])
+#     with gr.Row():
+#         with gr.Column():
+#             gr.Markdown("## Documentation Links")
+#             docs = gr.Markdown(default[3])
+#     inp.change(fn=change, inputs=[inp, "big_model_inference"], outputs=[out, feature, explanation, docs])
+
+
+# def notebook_launcher():
+#     inp = gr.Radio(
+#         ["Colab GPU", "Colab TPU", "Kaggle GPU", "Kaggle Multi GPU", "Kaggle TPU", "Multi GPU VMs"],
+#         label="Select a feature you would like to integrate",
+#         value="Basic",
+#     )
+#     with gr.Row():
+#         with gr.Column():
+#             feature = gr.Markdown("## Accelerate Code")
+#             out = gr.Markdown(default[0])
+#     with gr.Row():
+#         with gr.Column():
+#             gr.Markdown(default[1])
+#             explanation = gr.Markdown(default[2])
+#     with gr.Row():
+#         with gr.Column():
+#             gr.Markdown("## Documentation Links")
+#             docs = gr.Markdown(default[3])
+#     inp.change(fn=change, inputs=[inp, "notebook_launcher"], outputs=[out, feature, explanation, docs])
+
+
+with gr.Blocks() as demo:
+
+    with gr.Tabs():
+        with gr.TabItem("Simplify your code and improve efficieny"):
+            textbox = gr.Textbox(label="tab_name", visible=False, value="base")
+            base_features(textbox)
+        with gr.TabItem("Large Scale Training"):
+            textbox = gr.Textbox(label="tab_name", visible=False, value="large_scale_training")
+            large_scale_training(textbox)
+        with gr.TabItem("Big Model Inference"):
+            # big_model_inference()
+            pass
+        with gr.TabItem("Notebook Launcher Intergation"):
+            # notebook_launcher()
+            pass
+
+
+demo.launch()

src/markup.py

@@ -17,6 +17,7 @@ from template import get_filename
 _remove_color = "rgb(103,6,12)"
 _addition_color = "rgb(6,103,12)"
 
+
 def mark_text(text, add=True):
     """Marks text with a highlight color for addition or removal.
 
@@ -35,7 +36,8 @@ def mark_text(text, add=True):
         color = _remove_color
     return f'<mark style="background-color:{color}!important;color:white!important">{text}</mark>'
 
-def highlight(code:str):
+
+def highlight(code: str):
     """Takes in code and returns the respective highlighted code sample.
 
     Args:
@@ -43,7 +45,7 @@ def highlight(code:str):
             Code from a file.
     """
     lines = code.split("\n")
-    for i,line in enumerate(lines):
+    for i, line in enumerate(lines):
         if line.startswith("-"):
             lines[i] = "- " + line[1:]
             lines[i] = mark_text(lines[i], False)
@@ -54,12 +56,12 @@ def highlight(code:str):
             lines[i] = "  " + line
     return "\n".join(lines).rstrip()
 
-def get_text(option):
+
+def get_text(option, tab):
     """
     Reads in an option and returns the code, explanation, and documentation links
     """
-    filename = option.lower().replace(' ', '_')
-    with open(get_filename(filename)) as f:
+    filename = option.lower().replace(" ", "_")
+    with open(get_filename(tab, filename)) as f:
         output = f.read()
-    code, explanation, doclink = output.split("##\n")[1:]
-    return code, explanation, doclink
+    return output.split("##\n")[1:]

src/template.py

@@ -15,17 +15,16 @@ import os
 
 TEMPLATES = ["initial", "initial_with_metrics", "accelerate"]
 
-def get_filename(template: str) -> str:
+
+def get_filename(tab: str, template: str) -> str:
     """
     Takes an template and returns the respective filename relative to the cwd.
     """
-    return os.path.join(os.getcwd(), "code_samples", template)
+    return os.path.join(os.getcwd(), "code_samples", tab, template)
+
 
 def get_templates() -> dict:
     """
     Returns a dictionary of template type to code content
     """
-    return {
-        template: open(get_filename(template)).read()
-        for template in TEMPLATES
-    }
+    return {template: open(get_filename("base", template)).read() for template in TEMPLATES}