set fsdp and bf16 don't save memory

System Info

transformers version: 4.28.0
Platform: Linux-5.4.0-122-generic-x86_64-with-glibc2.31
Python version: 3.9.12
Huggingface_hub version: 0.13.4
Safetensors version: not installed
PyTorch version (GPU?): 1.13.1+cu117 (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?

@ArthurZucker @sgu

Information

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

Tasks

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

Reproduction

download the dataset
```
lang = "Python"
```

import subprocess subprocess.call(["wget", f"https://s3.amazonaws.com/code-search-net/CodeSearchNet/v2/{lang}.zip"]) subprocess.call(["unzip", f"/content/{lang}.zip"])

!mkdir "log" log_dir = "/content/log" !mkdir "data" data_dir = "/content/data" !mkdir "model" model_dir = "/content/model" !mkdir "tokenizer" tokenizer_dir = "/content/tokenizer"


2. data preprocess

import os import json import torch from pathlib import Path from transformers import (Trainer, pipeline, RobertaConfig, TrainingArguments, RobertaForMaskedLM, RobertaTokenizerFast, LineByLineTextDataset, DataCollatorForLanguageModeling)

from tokenizers import ByteLevelBPETokenizer from tokenizers.processors import BertProcessing from tokenizers.implementations import ByteLevelBPETokenizer

def prepare_text(dir_path): for path in os.listdir(dir_path): os.system(f"gunzip -k {dir_path}/{path}")

texts = "" for path in os.listdir(dir_path): if path.endswith(".jsonl"): with open(dir_path + "/" + path, 'r') as f: sample_file = f.readlines() for sample in sample_file: obj = json.loads(sample) texts += obj["original_string"].replace("\n", "").replace("\t", "") + "\n" return texts

train1_texts = prepare_text(f"/content/{lang}/final/jsonl/train") train2_texts = prepare_text(f"/content/{lang}/final/jsonl/valid") train_texts = train1_texts + "\n" + train2_texts valid_texts = prepare_text(f"/content/{lang}/final/jsonl/test")

for path, text in zip(["train_texts.txt", "valid_texts.txt"], [train_texts, valid_texts]): with open(f"{data_dir}/{path}","w") as f: f.write(text)


3. Train a tokenizer

paths = [str(x) for x in Path(f"{data_dir}/").glob("*/.txt")] tokenizer = ByteLevelBPETokenizer()

tokenizer.train(files=paths, vocab_size=52_000, min_frequency=2, special_tokens=[ "~~", "", "~~", "", "", ])

tokenizer.save_model(tokenizer_dir)

tokenizer = ByteLevelBPETokenizer( "tokenizer/vocab.json", "tokenizer/merges.txt", )

tokenizer._tokenizer.post_processor = BertProcessing( ("", tokenizer.token_to_id("")), ("~~", tokenizer.token_to_id("~~")), ) tokenizer.enable_truncation(max_length=512)~~~~

4. Build model

config = RobertaConfig( vocab_size=52_000, max_position_embeddings=514, num_attention_heads=12, num_hidden_layers=6, type_vocab_size=1, )

tokenizer = RobertaTokenizerFast.from_pretrained(tokenizer_dir, max_len=512)

model = RobertaForMaskedLM(config=config) model.num_parameters()

train_dataset = LineByLineTextDataset( tokenizer=tokenizer, file_path=f"{data_dir}/train_texts.txt", block_size=128, )

test_dataset = LineByLineTextDataset( tokenizer=tokenizer, file_path=f"{data_dir}/valid_texts.txt", block_size=128, )

data_collator = DataCollatorForLanguageModeling( tokenizer=tokenizer, mlm=True, mlm_probability=0.15 )

training_args = TrainingArguments( output_dir=model_dir, overwrite_output_dir=True, num_train_epochs=4, per_gpu_train_batch_size=64, save_steps=5000, do_eval=True, logging_dir=log_dir, )

trainer = Trainer( model=model, args=training_args, data_collator=data_collator, train_dataset=train_dataset, eval_dataset = test_dataset )

trainer.train()

trainer.save_model(model_dir)

tokenizer.save_pretrained(tokenizer_dir)

### Expected behavior before set fsdp and bf16:

training_args = TrainingArguments( output_dir=model_dir, overwrite_output_dir=True, num_train_epochs=4, per_gpu_train_batch_size=64, save_steps=5000, do_eval=True, logging_dir=log_dir, )

<img width="417" alt="Snipaste_2023-04-18_15-42-22" src="https://user-images.githubusercontent.com/41561936/232707188-2579965b-92fd-4ba6-87de-b82ca948ec54.png"> after set fsdp and bf16:

training_args = TrainingArguments( output_dir=model_dir, overwrite_output_dir=True, num_train_epochs=4, per_gpu_train_batch_size=64, save_steps=5000, do_eval=True, logging_dir=log_dir, fsdp=True, bf16=True, )

<img width="415" alt="Snipaste_2023-04-18_15-42-45" src="https://user-images.githubusercontent.com/41561936/232707483-2b89c658-172d-4a23-a7fc-fe40cd1dfe83.png"> The memory usage is not much different and does not achieve the desired effect. Why? I also try to set `per_gpu_train_batch_size=4` when `fsdp=True, bf16=True`: <img width="426" alt="Snipaste_2023-04-18_15-49-23" src="https://user-images.githubusercontent.com/41561936/232708818-efa676d9-4e6b-440a-b0e0-e66e54026da5.png"> Compared with the results of the previous set of experiments, the increase of memory usage is much greater than the increase of batch size. Why?

huggingface / transformers