pref: 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 default_data_collator, DataCollatorWithPadding
 
 from bsmetadata.metadata_utils import add_metadata_and_chunk_examples
 
@@ -124,15 +125,22 @@ 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 default_data_collator, DataCollatorWithPadding
 
 
 logger = logging.getLogger(__name__)
@@ -157,15 +158,22 @@ 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/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)