Adding FSDP Support to Training Library

.pylintrc

@@ -472,7 +472,8 @@ disable=raw-checker-failed,
         consider-using-generator,
         broad-exception-caught,
         super-init-not-called,
-        duplicate-code
+        duplicate-code,
+        too-many-positional-arguments
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

requirements.txt

@@ -6,6 +6,7 @@ py-cpuinfo
 # replace custom pytorch images with the 2.3.0
 torch>=2.3.0a0
 transformers>=4.41.2
+accelerate>=0.34.2
 datasets>=2.15.0
 numba
 # Note: numpy ranges copied from instructlab/instructlab

src/instructlab/training/__init__.py

@@ -7,15 +7,21 @@
     "TorchrunArgs",
     "TrainingArgs",
     "run_training",
+    "FSDPOptions",
+    "ShardingStrategies",
+    "DistributedTrainingBackend",
 )
 
 # Local
 from .config import (
     DataProcessArgs,
     DeepSpeedOffloadStrategy,
     DeepSpeedOptions,
+    DistributedTrainingBackend,
+    FSDPOptions,
     LoraOptions,
     QuantizeDataType,
+    ShardingStrategies,
     TorchrunArgs,
     TrainingArgs,
 )

src/instructlab/training/chat_templates/ibm_generic_tmpl.py

@@ -1,14 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # First Party
-from instructlab.training.tokenizer_utils import SpecialTokens
+from instructlab.training.tokenizer_utils import SpecialTokens, TokenInfo
 
 SPECIAL_TOKENS = SpecialTokens(
-    system="<|system|>",
-    user="<|user|>",
-    assistant="<|assistant|>",
-    eos="<|endoftext|>",
-    pad="<|pad|>",
+    system=TokenInfo("<|system|>", add_to_tokenizer=True),
+    user=TokenInfo("<|user|>", add_to_tokenizer=True),
+    assistant=TokenInfo("<|assistant|>", add_to_tokenizer=True),
+    eos=TokenInfo("<|endoftext|>", add_to_tokenizer=True),
+    pad=TokenInfo("<|pad|>", add_to_tokenizer=True),
+    bos=TokenInfo("<|begginingoftext|>", add_to_tokenizer=True),
 )
 
 CHAT_TEMPLATE = (

src/instructlab/training/chat_templates/mistral_tmpl.py

@@ -1,24 +1,39 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # First Party
-from instructlab.training.tokenizer_utils import SpecialTokens
+from instructlab.training.tokenizer_utils import SpecialTokens, TokenInfo
 
 SPECIAL_TOKENS = SpecialTokens(
-    bos="<s>",
-    eos="</s>",
-    user="[INST]",
-    assistant="[/INST]",
+    bos=TokenInfo("<s>", add_to_tokenizer=True),
+    eos=TokenInfo("</s>", add_to_tokenizer=True),
+    user=TokenInfo("[INST]", add_to_tokenizer=False),
+    assistant=TokenInfo("[/INST]", add_to_tokenizer=False),
 )
 
 CHAT_TEMPLATE = (
+    "{%- if messages[0]['role'] == 'system' %}"
+    "{%- set system_message = messages[0]['content'] %}"
+    "{%- set loop_messages = messages[1:] %}"
+    "{%- else %}"
+    "{%- set loop_messages = messages %}"
+    "{%- endif %}"
     "{{ '<s>' }}"
-    "{% for message in messages %}"
-    "{% if message['role'] == 'pretraining' %}"
-    "{{'<|pretrain|>' + message['content'] + '</s>' + '<|/pretrain|>'}}"
-    "{% elif message['role'] == 'user' %}"
-    "{{ '[INST] ' + message['content'] + ' [/INST]' }}"
-    "{% elif message['role'] == 'assistant' %}"
-    "{{ message['content'] + '</s>'}}"
-    "{% endif %}"
-    "{% endfor %}"
+    "{%- for message in loop_messages %}"
+    "{%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}"
+    "{{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}"
+    "{%- endif %}"
+    "{%- if message['role'] == 'user' %}"
+    "{%- if loop.first and system_message is defined %}"
+    "{{- ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST]' }}"
+    "{%- else %}"
+    "{{- ' [INST] ' + message['content'] + ' [/INST]' }}"
+    "{%- endif %}"
+    "{%- elif message['role'] == 'pretraining' %}"
+    "{{- '<|pretrain|>' + message['content'] + '</s>' + '<|/pretrain|>' }}"
+    "{%- elif message['role'] == 'assistant' %}"
+    "{{- ' ' + message['content'] + '</s>'}}"
+    "{%- else %}"
+    "{{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}"
+    "{%- endif %}"
+    "{%- endfor %}"
 )

src/instructlab/training/config.py

@@ -28,6 +28,12 @@ class DeepSpeedOffloadStrategy(Enum):
     NONE = None
 
 
+# public API
+class DistributedTrainingBackend(Enum):
+    FSDP: str = "fsdp"
+    DEEPSPEED: str = "deepspeed"
+
+
 # public API
 class QuantizeDataType(Enum):
     """
@@ -111,6 +117,24 @@ class DeepSpeedOptions(BaseModel):
     save_samples: int | None = None
 
 
+# public API
+class ShardingStrategies(Enum):
+    FULL_SHARD = "FULL_SHARD"
+    SHARD_GRAD_OP = "SHARD_GRAD_OP"
+    NO_SHARD = "NO_SHARD"
+    HYBRID_SHARD = "HYBRID_SHARD"
+
+
+# public API
+class FSDPOptions(BaseModel):
+    """
+    Represents the options for configuring FSDP which are exposed by the Training Library
+    """
+
+    cpu_offload_params: Optional[bool] = False
+    sharding_strategy: ShardingStrategies = ShardingStrategies.SHARD_GRAD_OP
+
+
 # public API
 class TrainingArgs(BaseModel):
     """
@@ -157,6 +181,14 @@ class TrainingArgs(BaseModel):
             cpu_offload_optimizer_pin_memory=False,
         )
     )
+    fsdp_options: FSDPOptions = Field(
+        default_factory=lambda: FSDPOptions(
+            cpu_offload_params=False, sharding_strategy=ShardingStrategies.SHARD_GRAD_OP
+        )
+    )
+    distributed_training_backend: DistributedTrainingBackend = (
+        DistributedTrainingBackend.DEEPSPEED
+    )
 
     disable_flash_attn: Optional[bool] = False