wip

abhi-mosaic · abhi-mosaic · commit c19fc86da9cd · 2022-01-12T01:20:15.000Z
diff --git a/composer/datasets/streaming_lm_datasets.py b/composer/datasets/streaming_lm_datasets.py
@@ -14,9 +14,9 @@
 import yahp as hp
 from transformers.testing_utils import CaptureLogger
 
-from composer.core.types import Batch
+from composer.core.types import Batch, DataSpec
 from composer.datasets.dataloader import DataloaderHparams
-from composer.datasets.hparams import DataloaderSpec, DatasetHparams
+from composer.datasets.hparams import DatasetHparams
 from composer.utils import dist
 from composer.utils.data import get_subset_dataset
 
@@ -73,18 +73,18 @@ def _load_dataset(self):
                                      split=self.split,
                                      streaming=True)
 
-    def _get_approx_num_samples(self):
+    def _get_approx_num_samples_per_device(self):
         try:
             if self.max_samples > 0:
-                return self.max_samples
+                return self.max_samples // dist.get_world_size()
             else:
                 n_shards, samples_per_shard = CACHED_DATASET_SIZES[self.dataset_name][self.dataset_config_name][self.split]
                 n_shards = self.max_shards if self.max_shards > 0 else n_shards
-                return n_shards * samples_per_shard
+                return n_shards * samples_per_shard // dist.get_world_size()
         except:
             raise NotImplementedError
 
-    def _get_approx_num_tokens(self):
+    def _get_approx_num_tokens_per_device(self):
         return 1e12
 
     def _subsample(self, device_offset, text_batch):
@@ -166,7 +166,7 @@ def _group_tokens(self, token_batch):
         else:
             raise ValueError(f"Unknown group_method: '{group_method}'")
 
-    def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHparams) -> DataloaderSpec:
+    def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHparams) -> DataSpec:
         assert dataloader_hparams.num_workers == 1, "LM Streaming Dataloader only supports num_workers=1"
 
         try:
@@ -209,13 +209,12 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
                 batch_size=token_sample_batch_size,
             )
 
-        # Maybe limit the number of post-processed samples
-        if self.max_samples > 0:
-            token_dataset = token_dataset.take(self.max_samples // dist.get_world_size())
-
-        # Add approx num samples and create a SizedIterableDataset
-        sized_iterable_dataset = SizedIterableDataset(token_dataset, self._get_approx_num_samples())
+        # Limit the number of post-processed samples
+        num_samples_per_device = self._get_approx_num_samples_per_device()
+        token_dataset = token_dataset.take(num_samples_per_device)
 
+        # HACK: create a SizedIterableDataset
+        sized_iterable_dataset = SizedIterableDataset(token_dataset, num_samples_per_device)
 
         # Get collate_fn
         if self.tokenizer_name in ["gpt2"]:
@@ -225,25 +224,25 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
             collate_fn = transformers.DataCollatorForLanguageModeling(tokenizer=self.tokenizer,
                                                                   mlm=self.use_masked_lm,
                                                                   mlm_probability=self.mlm_probability)
-        # Return DataloaderSpec
-        return DataloaderSpec(dataloader=dataloader_hparams.initialize_object(
+        # Return DataSpec
+        return DataSpec(dataloader=dataloader_hparams.initialize_object(
             dataset=sized_iterable_dataset,
             batch_size=batch_size,
             sampler=None,
             drop_last=self.drop_last,
             collate_fn=collate_fn,
         ),
-                              split_fn=_split_dict_fn)
+                              split_batch=_split_dict_fn)
 
 
 class SizedIterableDataset(torch.utils.data.IterableDataset):
 
-    def __init__(self, hf_iterable_dataset, num_samples):
+    def __init__(self, hf_iterable_dataset, num_samples_per_device):
         self.hf_iterable_dataset = hf_iterable_dataset
-        self.num_samples = num_samples
+        self.num_samples_per_device = num_samples_per_device
 
     def __iter__(self):
         return iter(self.hf_iterable_dataset)
 
     def __len__(self):
-        return self.num_samples
+        return self.num_samples_per_device
diff --git a/composer/trainer/deepspeed.py b/composer/trainer/deepspeed.py
@@ -83,7 +83,7 @@ def initialize_object(self, state: State, grad_clip_norm: Optional[float]):
         elif state.precision == Precision.FP16:
             deepspeed_config["fp16"] = {
                 "enabled": True,
-                "initial_scale_power": 16,
+                "initial_scale_power": 0,
                 "loss_scale_window": 2000,
             }
 
diff --git a/composer/trainer/trainer_hparams.py b/composer/trainer/trainer_hparams.py
@@ -204,11 +204,10 @@ class TrainerHparams(hp.Hparams):
     def validate(self):
         super().validate()
 
-        if self.deepspeed is not None:
-
+        if self.deepspeed is None:
             if self.precision == Precision.FP16:
                 raise ValueError("FP16 precision is only supported when training with DeepSpeed.")
-
+        else:
             if isinstance(self.device, CPUDeviceHparams):
                 raise ValueError("Training on CPUs is not supported with DeepSpeed.")
 
diff --git a/composer/utils/_time_conversion.py b/composer/utils/_time_conversion.py
@@ -42,9 +42,9 @@ def convert(
             converting to or from :attr:`TimeUnit.TOKEN`.
         max_training_duration (str or Time, optional): The total training duration. Required only
             if converting to or from :attr:`TimeUnit.DURATION`.
-    
+
     Raises:
-        ValueError: If it is not possible to perform the conversion. 
+        ValueError: If it is not possible to perform the conversion.
 
     Returns:
         Time: The time, in the specified ``unit``.
@@ -76,12 +76,13 @@ def convert(
                                                 dataset_num_tokens=dataset_num_tokens)
             return _convert_to_duration(time_in_max_duration_unit, max_training_duration=max_training_duration)
         else:
-            converted_time = _convert_from_duration(time, max_training_duration=max_training_duration)
-            return convert(converted_time,
-                           unit,
-                           steps_per_epoch=steps_per_epoch,
-                           samples_per_epoch=samples_per_epoch,
-                           dataset_num_tokens=dataset_num_tokens)
+            max_training_duration_in_unit = convert(max_training_duration,
+                                            unit,
+                                            steps_per_epoch=steps_per_epoch,
+                                            samples_per_epoch=samples_per_epoch,
+                                            dataset_num_tokens=dataset_num_tokens)
+            converted_time = _convert_from_duration(time, max_training_duration=max_training_duration_in_unit)
+            return converted_time
 
     if time.unit == TimeUnit.EPOCH:
         if unit == TimeUnit.BATCH:
@@ -260,7 +261,7 @@ def _convert_sample_to_batch(
         time (Time): The time
         steps_per_epoch (int): The number of optimization steps per epoch.
         samples_per_epoch (int): The number of samples per epoch.
-    
+
     Raises:
         RuntimeError: Raised if ``time.unit != TimeUnit.SAMPLE``
 
diff --git a/composer/yamls/models/gpt3_1,3b.yaml b/composer/yamls/models/gpt3_1,3b.yaml
@@ -0,0 +1,116 @@
+# GPT3-125m with streaming C4 dataset
+
+train_dataset:
+  streaming_lm:
+    dataset_name: c4
+    dataset_config_name: en
+    split: train
+    max_shards: -1
+    max_samples: 5120  # 512sa * 10ba
+    max_seq_len: 2048
+    group_method: concat
+    tokenizer_name: gpt2
+    use_masked_lm: false
+    seed: 17
+    shuffle: true
+    drop_last: true
+val_dataset:
+  streaming_lm:
+    dataset_name: c4
+    dataset_config_name: en
+    split: validation
+    max_shards: -1
+    max_samples: 100
+    max_seq_len: 2048
+    group_method: concat
+    tokenizer_name: gpt2
+    use_masked_lm: false
+    seed: 17
+    shuffle: false
+    drop_last: true
+
+model:
+  gpt2:
+    use_pretrained: false
+    tokenizer_name: gpt2
+    model_config:
+      activation_function: gelu_new
+      architectures:
+        - GPT2LMHeadModel
+      attn_pdrop: 0.1
+      bos_token_id: 50256
+      embd_pdrop: 0.1
+      eos_token_id: 50256
+      initializer_range: 0.02
+      layer_norm_epsilon: 1.0e-05
+      model_type: gpt2
+      n_embd: 2048
+      n_head: 16
+      n_inner: 8192
+      n_layer: 24
+      n_positions: 2048
+      resid_pdrop: 0.1
+      scale_attn_weights: true
+      summary_activation: null
+      summary_first_dropout: 0.1
+      summary_proj_to_labels: true
+      summary_type: cls_index
+      summary_use_proj: true
+      task_specific_params:
+        text-generation:
+          do_sample: true
+          max_length: 50
+      transformers_version: 4.11.0.dev0
+      use_cache: true
+      vocab_size: 50257
+optimizer:
+  decoupled_adamw:
+    lr: 2.0e-4
+    betas:
+      - 0.9
+      - 0.95
+    eps: 1.0e-08
+    weight_decay: 0.0
+schedulers:
+  - warmup:
+      warmup_method: linear
+      warmup_iters: 0.2dur
+      warmup_factor: 0
+      interval: batch
+  - linear_decay:
+      start_factor: 1.0
+      end_factor: 0.0
+      total_iters: 0.8dur
+      interval: batch
+      verbose: false
+loggers:
+  - file:
+      log_level: batch
+      filename: stdout
+      buffer_size: 1
+      flush_every_n_batches: 100
+      every_n_batches: 1
+      every_n_epochs: 1
+max_duration: 1ep
+train_batch_size: 512
+eval_batch_size: 8 # use micro_bs_per_gpu = 1 to accomodate 10GB limit
+seed: 17
+device:
+  gpu: {}
+deepspeed:
+  zero_stage: 0
+  # optimizer_offload: true
+  # parameter_offload: true
+  # overlap_comm: false
+  # gradient_checkpointing: false
+dataloader:
+  pin_memory: true
+  persistent_workers: true
+  num_workers: 1
+  timeout: 0
+  prefetch_factor: 2
+precision: fp16
+grad_clip_norm: 1.0
+grad_accum: 1
+validate_every_n_batches: 3
+validate_every_n_epochs: 1
diff --git a/composer/yamls/models/gpt3_125m.yaml b/composer/yamls/models/gpt3_125m.yaml
@@ -1,13 +1,13 @@
-# GPT2-125m with streaming C4 dataset
+# GPT3-125m with streaming C4 dataset
 
 train_dataset:
   streaming_lm:
     dataset_name: c4
     dataset_config_name: en
     split: train
     max_shards: -1
-    max_samples: 7168000
-    max_seq_len: 1024
+    max_samples: 2560  # 256sa * 20ba
+    max_seq_len: 2048
     group_method: concat
     tokenizer_name: gpt2
     use_masked_lm: false
@@ -20,8 +20,8 @@ val_dataset:
     dataset_config_name: en
     split: validation
     max_shards: -1
-    max_samples: 128000
-    max_seq_len: 1024
+    max_samples: 100
+    max_seq_len: 2048
     group_method: concat
     tokenizer_name: gpt2
     use_masked_lm: false
@@ -44,12 +44,11 @@ model:
       initializer_range: 0.02
       layer_norm_epsilon: 1.0e-05
       model_type: gpt2
-      n_ctx: 1024
       n_embd: 768
       n_head: 12
       n_inner: 3072
       n_layer: 12
-      n_positions: 1024
+      n_positions: 2048
       resid_pdrop: 0.1
       scale_attn_weights: true
       summary_activation: null
@@ -65,34 +64,35 @@ model:
       use_cache: true
       vocab_size: 50257
 optimizer:
-  adamw:
+  decoupled_adamw:
     lr: 6.0e-4
     betas:
       - 0.9
-      - 0.999
+      - 0.95
     eps: 1.0e-08
     weight_decay: 0.0
 schedulers:
   - warmup:
       warmup_method: linear
+      warmup_iters: 0.2dur
       warmup_factor: 0
-      interval: step
-      warmup_iters: 140ba
-  - cosine_decay:
-      interval: step
-      eta_min: 0
+      interval: batch
+  - linear_decay:
+      start_factor: 1.0
+      end_factor: 0.0
+      total_iters: 0.8dur
+      interval: batch
       verbose: false
-      T_max: 13860ba
 loggers:
   - file:
       log_level: batch
       filename: stdout
       buffer_size: 1
       flush_every_n_batches: 100
-      every_n_batches: 100
+      every_n_batches: 1
       every_n_epochs: 1
-max_epochs: 1
-train_batch_size: 512
+max_duration: 1ep
+train_batch_size: 256
 eval_batch_size: 8 # use micro_bs_per_gpu = 1 to accomodate 10GB limit
 seed: 17
 device:
@@ -105,6 +105,6 @@ dataloader:
   prefetch_factor: 2
 precision: amp
 grad_clip_norm: 1.0
-grad_accum: 22
-validate_every_n_batches: 1000
+grad_accum: 1
+validate_every_n_batches: 3
 validate_every_n_epochs: 1

Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,7 @@ def initialize_object(self, state: State, grad_clip_norm: Optional[float]):`
`83`	`83`	`elif state.precision == Precision.FP16:`
`84`	`84`	`deepspeed_config["fp16"] = {`
`85`	`85`	`"enabled": True,`
`86`		`- "initial_scale_power": 16,`
	`86`	`+ "initial_scale_power": 0,`
`87`	`87`	`"loss_scale_window": 2000,`
`88`	`88`	`}`
`89`	`89`