add falcon model, and lima linear dropout increase
diff --git a/model/model_training/configs/config.yaml b/model/model_training/configs/config.yaml
index 97a72c951e..8266025a77 100644
--- a/model/model_training/configs/config.yaml
+++ b/model/model_training/configs/config.yaml
@@ -83,6 +83,7 @@ defaults:
   system_add_length: false
   per_digit_tokens: false
   is_reward_model: false
+  residual_dropout_lima: false
   deepspeed_config: configs/zero_config.json
 
 use_system_tag:
@@ -224,6 +225,38 @@ oasst_export_latin_cyrillic:
   sort_by_length: false
   use_custom_sampler: false
 
+oasst-top1:
+  save_strategy: steps # epoch seems not to work, gets stuck with DS 0.9.1
+  save_steps: 600
+  datasets:
+    - oasst_export:
+        lang: "bg,ca,cs,da,de,en,es,fr,hr,hu,it,nl,pl,pt,ro,ru,sl,sr,sv,uk" # sft-8.0
+        input_file_path: 2023-05-06_OASST_labels.jsonl.gz
+        val_split: 0.05
+        top_k: 1
+
+falcon-7b:
+  dtype: bf16
+  log_dir: "llama_log_7b"
+  learning_rate: 1e-5
+  model_name: "tiiuae/falcon-7b"
+  deepspeed_config: configs/zero_config_falcon.json
+  output_dir: falcon
+  weight_decay: 0.0
+  max_length: 2048
+  warmup_steps: 100
+  gradient_checkpointing: true
+  gradient_accumulation_steps: 2
+  per_device_train_batch_size: 1
+  per_device_eval_batch_size: 1
+  eval_steps: 100
+  save_steps: 500
+  num_train_epochs: 8
+  save_total_limit: 4
+  use_flash_attention: false
+  residual_dropout: 0.3
+  residual_dropout_lima: true
+
 llama-7b:
   dtype: fp16
   log_dir: "llama_log_7b"
diff --git a/model/model_training/configs/zero_config_falcon.json b/model/model_training/configs/zero_config_falcon.json
new file mode 100644
index 0000000000..c9ba7902a7
--- /dev/null
+++ b/model/model_training/configs/zero_config_falcon.json
@@ -0,0 +1,29 @@
+{
+  "bf16": {
+    "enabled": "auto"
+  },
+  "scheduler": {
+    "type": "WarmupDecayLR",
+    "params": {
+      "warmup_min_lr": "auto",
+      "warmup_max_lr": "auto",
+      "warmup_num_steps": "auto",
+      "total_num_steps": "auto"
+    }
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 1e9,
+    "overlap_comm": false,
+    "reduce_scatter": true,
+    "reduce_bucket_size": 1e9,
+    "contiguous_gradients": true
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "steps_per_print": 2000,
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}
diff --git a/model/model_training/models/__init__.py b/model/model_training/models/__init__.py
index 73ece77f56..50c6ca7eb4 100644
--- a/model/model_training/models/__init__.py
+++ b/model/model_training/models/__init__.py
@@ -35,5 +35,7 @@ def get_specific_model(
     elif seq2seqmodel:
         model = transformers.AutoModelForSeq2SeqLM.from_pretrained(model_name, cache_dir=cache_dir, **kwargs)
     else:
+        if "falcon" in model_name:
+            kwargs["trust_remote_code"] = True
         model = transformers.AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir, **kwargs)
     return model
diff --git a/model/model_training/models/patching.py b/model/model_training/models/patching.py
index 4a17f5c60c..c8757beb8f 100644
--- a/model/model_training/models/patching.py
+++ b/model/model_training/models/patching.py
@@ -73,6 +73,7 @@ def patch_model(
     resid_pdrop: Optional[float] = 0.1,
     flash_attention: bool = True,
     patch_unsupported: bool = False,
+    residual_dropout_lima: bool = False,
 ):
     """
     Helper function for patching HF language models.
@@ -104,7 +105,10 @@ def patch_model(
     if not flash_attention and (resid_pdrop is None or resid_pdrop == 0.0):
         return
 
-    if not any(isinstance(model, model_class) for model_class in SUPPORTED_MODELS):
+    if (
+        not any(isinstance(model, model_class) for model_class in SUPPORTED_MODELS)
+        and model.__class__.__name__ != "RWForCausalLM"
+    ):
         if not flash_attention and (resid_pdrop is None or resid_pdrop == 0.0):
             return  # nothing to patch
 
@@ -141,6 +145,9 @@ def patch_model(
                 "--use_flash_attention=false  --no-residual_dropout"
             )
 
+    if model.__class__.__name__ == "RWForCausalLM":
+        model = model.base_model
+
     attention_key_lookup = {
         GPTNeoXModel: "attention",
         GPTNeoXRewardModel: "attention",
@@ -151,13 +158,21 @@ def patch_model(
         GPTNeoXRewardModel: "mlp",
         LlamaModel: "mlp",
     }
-    attention_key = attention_key_lookup.get(model.__class__, "attention")
-    mlp_key = mlp_key_lookup.get(model.__class__, "mlp")
-
-    for layer in model.layers:
+    if model.__class__.__name__ == "RWModel":
+        layers = model.h
+        attention_key = "self_attention"
+        mlp_key = "mlp"
+    else:
+        layers = model.layers
+        attention_key = attention_key_lookup.get(model.__class__, "attention")
+        mlp_key = mlp_key_lookup.get(model.__class__, "mlp")
+    num_layers = len(layers)
+    resid_pdrop_last_layer = resid_pdrop
+    for i, layer in enumerate(layers):
         if flash_attention:
             add_flash_attn(getattr(layer, attention_key), causal=True)
-
+        if residual_dropout_lima:
+            resid_pdrop = i / (num_layers - 1) * resid_pdrop_last_layer
         if resid_pdrop is not None and resid_pdrop > 0:
             add_dropout(getattr(layer, attention_key), _patched_attn_forward, resid_pdrop)
             add_dropout(getattr(layer, mlp_key), _patched_mlp_forward, resid_pdrop)
diff --git a/model/model_training/utils/utils.py b/model/model_training/utils/utils.py
index 8e72970e09..bb04919f18 100644
--- a/model/model_training/utils/utils.py
+++ b/model/model_training/utils/utils.py
@@ -181,6 +181,9 @@ class TokenizerConfig(NamedTuple):
     "deberta-v3": TokenizerConfig(special_tokens=SpecialTokens("[PAD]", "[SEP]", sep_token="[CLS]")),
     "bloom": TokenizerConfig(special_tokens=SpecialTokens("<pad>", "</s>", "<s>")),
     "electra": TokenizerConfig(special_tokens=SpecialTokens("[PAD]", "[SEP]", sep_token="[CLS]")),
+    "falcon": TokenizerConfig(
+        special_tokens=SpecialTokens("<|endoftext|>", "<|endoftext|>", sep_token="<|endoftext|>")
+    ),
 }
 
 
@@ -336,7 +339,12 @@ def get_model(conf, tokenizer, pad_vocab_size_to_multiple_of=16, check_freeze_la
     params = sum([p.numel() for p in model_parameters])
     print("Number of trainable parameters: {}M".format(int(params / 1e6)))
 
-    patch_model(model, resid_pdrop=conf.residual_dropout, flash_attention=conf.use_flash_attention)
+    patch_model(
+        model,
+        resid_pdrop=conf.residual_dropout,
+        flash_attention=conf.use_flash_attention,
+        residual_dropout_lima=conf.residual_dropout_lima,
+    )
 
     return model
 
