fix

README.md

@@ -12,7 +12,7 @@ Supported methods:
 
 ```python
 from transformers import AutoModelForSeq2SeqLM
-from pet import get_pet_config,get_pet_model
+from pet import get_pet_config, get_pet_model
 model_name_or_path = "bigscience/mt0-large"
 tokenizer_name_or_path = "bigscience/mt0-large"
 
@@ -28,9 +28,30 @@ pet_config = get_pet_config(config)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
 model = get_pet_model(model, pet_config)
 model.print_trainable_parameters()
-# output: 
+# output: trainable params: 2359296 || all params: 1231940608 || trainable%: 0.19151053100118282
 ```
 
+## PET + 🤗 Accelerate
+
+PET models work with 🤗 Accelerate out of the box. 
+For scaling to large models, you can leverage 🤗 Accelerate's PyTorch FSDP integration as shown below.
+PyTorch FSDP shards parameters, gradients and optimizer states across data parallel workers which enables
+large language models to fit on available hardware. 
+It also supports CPU offloading to further enable distributed training at scale. 
+The support for DeepSpeed ZeRO Stage-3 is currently in backlog.
+
+```python
+from pet.utils.other import fsdp_auto_wrap_policy
+
+...
+
+if accelerator.state.fsdp_plugin is not None:
+    accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model)
+
+model = accelerator.prepare(model)
+```
+
+
 ## Models support matrix
 
 ### Sequence Classification

src/pet/mapping.py

@@ -82,8 +82,8 @@ def _prepare_prompt_learning_config(pet_config, model_config):
             raise ValueError("Please specify `num_attention_heads` in `pet_config`")
         pet_config.num_attention_heads = num_attention_heads
 
-    if pet_config.encoder_hidden_size is None:
-        pet_config.encoder_hidden_size = token_dim
+    if getattr(pet_config, "encoder_hidden_size", None) is None:
+        setattr(pet_config, "encoder_hidden_size", token_dim)
 
     return pet_config