pref: truncate and pad to max length for TPU

Author: tianjianjiang

bsmetadata/experiments/with_metadata.py

@@ -1,9 +1,10 @@
 import functools
 import logging
 
+from accelerate import DistributedType
 from datasets import load_dataset
 from torch.utils.data import DataLoader
-from transformers import default_data_collator
+from transformers import DataCollatorWithPadding, default_data_collator
 
 from bsmetadata.metadata_utils import add_metadata_and_chunk_examples
 
@@ -124,15 +125,18 @@ def create_labels_column(examples):
     val_dataset = lm_datasets["validation"]
 
     # DataLoaders creation:
+    data_collator = default_data_collator
+    if args.distributed_type == DistributedType.TPU:
+        data_collator = DataCollatorWithPadding(tokenizer, padding="max_length", max_length=args.max_seq_len)
     train_dataloader = DataLoader(
         train_dataset,
         shuffle=True,
-        collate_fn=default_data_collator,
+        collate_fn=data_collator,
         batch_size=args.per_device_train_batch_size,
     )
     val_dataloader1 = DataLoader(
         val_dataset,
-        collate_fn=default_data_collator,
+        collate_fn=data_collator,
         batch_size=args.per_device_eval_batch_size,
     )
     return train_dataloader, {"val1": val_dataloader1}

bsmetadata/experiments/without_metadata.py

@@ -1,8 +1,9 @@
 import logging
 
+from accelerate import DistributedType
 from datasets import load_dataset
 from torch.utils.data import DataLoader
-from transformers import default_data_collator
+from transformers import DataCollatorWithPadding, default_data_collator
 
 
 logger = logging.getLogger(__name__)
@@ -94,7 +95,7 @@ def get_dataloaders(tokenizer, args):
     text_column_name = "text" if "text" in column_names else column_names[0]
 
     def tokenize_function(examples):
-        return tokenizer(examples[text_column_name])
+        return tokenizer(examples[text_column_name], truncation=True, max_length=args.max_seq_len)
 
     tokenized_datasets = raw_datasets.map(
         tokenize_function,
@@ -157,15 +158,18 @@ def group_texts(examples):
     val_dataset = lm_datasets["validation"]
 
     # DataLoaders creation:
+    data_collator = default_data_collator
+    if args.distributed_type == DistributedType.TPU:
+        data_collator = DataCollatorWithPadding(tokenizer, padding="max_length", max_length=args.max_seq_len)
     train_dataloader = DataLoader(
         train_dataset,
         shuffle=True,
-        collate_fn=default_data_collator,
+        collate_fn=data_collator,
         batch_size=args.per_device_train_batch_size,
     )
     val_dataloader1 = DataLoader(
         val_dataset,
-        collate_fn=default_data_collator,
+        collate_fn=data_collator,
         batch_size=args.per_device_eval_batch_size,
     )
     return train_dataloader, {"val1": val_dataloader1}

bsmetadata/input_pipeline.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass, field
 from typing import List, Optional
 
+from accelerate import DistributedType
+
 
 @dataclass
 class DataConfig:
@@ -62,6 +64,7 @@ class DataConfig:
     block_size: Optional[int] = field(
         default=None, metadata={"help": "Optional input sequence length after tokenization."}
     )
+    distributed_type: DistributedType = field(default=DistributedType.NO)
 
 
 def get_dataloaders(tokenizer, cfg: DataConfig):

bsmetadata/metadata_utils.py

@@ -51,7 +51,7 @@ def add_metadata_and_chunk_examples(
         if add_metadata:
             # Get the global metadata prefix that is prepended to each training example.
             global_metadata_prefix = create_global_metadata_prefix(example, cfg)
-            global_metadata_prefix_encoded = tokenizer.encode_plus(global_metadata_prefix).input_ids
+            global_metadata_prefix_encoded = tokenizer.encode_plus(global_metadata_prefix, truncation=True).input_ids
         else:
             global_metadata_prefix_encoded = []
 
@@ -66,7 +66,7 @@ def add_metadata_and_chunk_examples(
         if global_metadata_prefix_encoded:
             text_with_local_metadata = " " + text_with_local_metadata
         char_level_metadata_mask = [False] + char_level_metadata_mask
-        text_with_local_metadata_encoded = tokenizer.encode_plus(text_with_local_metadata)
+        text_with_local_metadata_encoded = tokenizer.encode_plus(text_with_local_metadata, truncation=True)
 
         def is_metadata(idx: int) -> bool:
             char_span = text_with_local_metadata_encoded.token_to_chars(idx)

bsmetadata/train.py

@@ -104,12 +104,14 @@ def loss_fn(batch, outputs, metadata_mask=None):
     return loss
 
 
-@hydra.main(config_name="config")
+@hydra.main(config_path=None, config_name="config")
 def main(args: CFG) -> None:
+    accelerator = Accelerator()
+    args.data_config.distributed_type = accelerator.distributed_type
+
     print(OmegaConf.to_yaml(args))
 
     set_seed(args.seed)
-    accelerator = Accelerator()
     is_local_main_process = accelerator.is_local_main_process
     tqdm = partial(original_tqdm, disable=not is_local_main_process, position=0)