Questions about Accelerate with FSDP
Author: bgub
Pull request: no

### System Info

- `transformers` version: 4.32.1
- Platform: Linux-3.10.0-1160.92.1.el7.x86_64-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: 0.22.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: YES
- Using distributed or parallel set-up in script?: YES

### Who can help?

@muellerzr, @pacman100 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

loop.py
```python
from transformers import LlamaForCausalLM, LlamaTokenizer, Trainer, TrainingArguments
from accelerate import Accelerator
import torch
import os
from get_data import train_dataset, eval_dataset, data_collator

accelerator = Accelerator()

SIZE = "7b"
MODEL_PATH = f"/mnt/models/llama2/hf/Llama-2-{SIZE}-hf"
NAME = f"llama2-{SIZE}-dolly-15k"
BATCH_SIZE = 8
NUM_EPOCHS = 3
OUTPUT_DIR = os.environ["SLURM_JOB_NAME"]

tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH, legacy=False)
tokenizer.pad_token = tokenizer.eos_token

model = LlamaForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.bfloat16)
model = accelerator.prepare(model)

training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=NUM_EPOCHS,
    learning_rate=2e-5,
    logging_steps=10,
    per_device_train_batch_size=BATCH_SIZE,
    remove_unused_columns=False,
    save_steps=1000,
    save_total_limit=1,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    data_collator=lambda x: data_collator(x, tokenizer),
)

trainer.train()
trainer.evaluate()

# TODO: model = accelerator.unwrap_model(model)

model.save_pretrained(f"/mnt/finetunes/{NAME}")
tokenizer.save_pretrained(f"/mnt/finetunes/{NAME}")
```

get_data.py
```python
from datasets import load_from_disk, disable_caching
disable_caching()

IGNORE_TOKEN = -100

#####################
# FORMAT DATA       #
#####################

template_context = """### Instruction:
{instruction}

### Context:
{context}

### Response:
"""

template_no_context = """### Instruction:
{instruction}

### Response:
"""

def data_to_string(data):

    instruction = data["instruction"]
    context = data["context"]
    response = data["response"]

    template = template_context if len(context) > 0 else template_no_context

    source = template.format(instruction=instruction, context=context)

    return {
        "source": source,
        "text": source + response, 
    }

original_dataset = load_from_disk("../datasets/databricks-dolly-15k")["train"] 

dataset = original_dataset.map(
    data_to_string
).remove_columns(
    original_dataset.column_names
).filter(
    lambda x: len(x["text"]) < 1000 # TODO: change to 4000
)

#####################
# SPLIT DATA        #
#####################

processed_dataset = dataset.train_test_split(test_size=0.1)

train_dataset = processed_dataset["train"]
eval_dataset = processed_dataset["test"]

#####################
# CREATE DATALOADER #
#####################


def data_collator(features, tokenizer):
    sources = [feature["source"] for feature in features]
    targets = [feature["text"] for feature in features]

    source_tokens = tokenizer(
        sources,
        return_tensors="pt",
        padding='longest',
        max_length=None,
    )

    target_tokens = tokenizer(
        targets,
        return_tensors="pt",
        padding='longest',
        max_length=None,
    )
        
    labels = target_tokens["input_ids"].clone()

    for i in range(len(labels)):
        source_len = source_tokens["attention_mask"][i].sum()
        
        labels[i, :source_len] = IGNORE_TOKEN

    res = {
        "input_ids": target_tokens["input_ids"],
        "attention_mask": target_tokens["attention_mask"],
        "labels": labels,
    }

    return res
```

accelerate_config.yaml
```yaml
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_offload_params: false
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 8
use_cpu: false
```

Command: `export SLURM_JOB_NAME="testingg" && accelerate launch --config_file ./accelerate_config.yaml loop.py`

### Expected behavior

## Questions:

1. When I run my training script, the progress bar shows only 1/8 of the batches I see when I use 1 process and `device_map="auto"`. Do the Trainer class and Accelerate coordinate to split batches across processes, sync gradients, and update the model? 

	The speed of the training (and suboptimal loss) make me fear that I'm only training the model on 1/8 of my data. If this is true, how can I handle distributed data correctly?

2. In my `data_collator` function, I return a special `labels` key with the prompt tokens masked to `-100`. Does `labels` get passed in to the model along with `input_ids`?

Comments:
1: Yes, we dispatch and split the batches. So a batch of 64 on one device when training on 8 GPUs, that batch of 64 is split 8 ways and sent to each GPU as one batch. See the docs here on this behavior: https://huggingface.co/docs/accelerate/concept_guides/performance#observed-batch-sizes

2: I *think* so, cc @amyeroberts?

Hello @nebrelbug, 

1. As Zach mentioned, in multi-gpu/multi-node setting, `per_device_train_batch_size*num_devices` will be the actual batch size. In your case, `per_device_train_batch_size*8` as you are running on 8 GPUs. So, the steps per epoch get reduced proportionally. In addition to the guide attached by Zach, note that when increasing batch size, one needs to increase the learning rate a bit. Also, rather than loss which is a factor of number of samples in batch which is huge when scaling across many GPUs, I would track the metric on eval dataset. 

2. Yes, they should be passed to the model. 

 

@muellerzr and @pacman100 thanks so much for the clarifications!

I didn't properly understand that how batch size works with a multi-GPU setup, or realize that I should increase learning rate! And it makes sense that loss isn't the best metric to use for evaluation.