update default test configs

EleutherAI · May 17, 2021 · a922b46 · a922b46
1 parent 130d9ad
commit a922b46
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/tests/model/test_model_train.py b/tests/model/test_model_train.py
@@ -16,12 +16,15 @@
  ["rmsnorm", "none", "swish"]],
  "pipe_parallel_size,model_parallel_size": [[0, 1], [1, 2], [0, 2]],
  "no_weight_tying": binary,
- "num_layers": [4],
- "attention_config": [[[["global"], "all"]], [[["local", "global"], "all"]], [[["sparse_variable", "global"], "all"]],
- [[["sparse_fixed", "global"], "all"]]],
+ "attention_config,num_layers": [[[[["global"], "all"]], 2], [[[["local", "global"], "all"]], 12], [[[["sparse_variable", "global"], "all"]], 12],
+ [[[["sparse_fixed", "global"], "all"]], 12]], # the sparse attention models need more layers to be stable
  "scaled_upper_triang_masked_softmax_fusion,bias_gelu_fusion": [[True, False], [False, True]],
  "checkpoint_activations": binary,
  "log_gradient_noise_scale": [True],
+ "sparsity_config": [{
+ "block": 16, # block size
+ "num_local_blocks": 32,
+ }]
 
 }
 

diff --git a/tests/test_configs/test_train_base.yml b/tests/test_configs/test_train_base.yml
@@ -7,7 +7,7 @@
 
  # model settings
  "num_layers": 2,
- "hidden_size": 96,
+ "hidden_size": 192,
  "num_attention_heads": 6,
  "seq_length": 1024,
  "max_position_embeddings": 1024,