# coding=utf-8 # Copyright 2021 The HuggingFace Team. All rights reserved. # Copyright (c) 2022 Graphcore Ltd. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import json import warnings from collections import defaultdict from functools import partial from tempfile import NamedTemporaryFile from typing import Any, Dict, List, Optional, Sequence, Union, get_type_hints import popart import poptorch import torch import typeguard from poptorch import Options, OutputMode from optimum.configuration_utils import BaseConfig from optimum.utils import logging # For container types check all items for type correctness # rather than just the first element typeguard._config.global_config.collection_check_strategy = typeguard.config.collection_check_strategy.ALL_ITEMS logger = logging.get_logger(__name__) IPU_CONFIG_NAME = "ipu_config.json" class IncompatibleIPUConfigError(ValueError): "Exception raised when IPU config values are invalid" "or are not compatible with a model" pass class IPUConfig(BaseConfig): """ Class for configuring PopArt and PyTorch for the IPU. Handles the conversion to `poptorch` options as well as configuration of the IPU-Pod type specialization. Args: seed (`int`, *optional*): Sets the seed for the random number generator on the IPU. auto_loss_scaling (`bool`, *optional*, defaults to `False`): If `True`, enables automatic loss scaling on the IPU. When using float16/half-precision values for activations, gradients, and weights, the loss value needs to be scaled by a constant factor to avoid underflows or overflows. This adjustment is known as loss scaling. This setting automatically sets a global loss scaling factor during training. **Note: This is an experimental feature and may not behave as expected.** executable_cache_dir (`str`, *optional*, defaults to `""`): Enables caching the compile executables to a directory. > Parameters for controlling the batch size replication_factor (`int`, *optional*, defaults to 1): The number of replicas for data-parallelism during training. It depends on the size of the pipeline as well as the number of IPUs available. For example: on a Pod16, with a 4-IPU pipeline, the replication_factor must be between 1 and 4. inference_replication_factor (`int`, *optional*, defaults to 1): The number of replicas for data-parallelism during inference. It depends on the size of the pipeline as well as the number of IPUs available. For example: on a Pod16, with a 4-IPU pipeline, the replication_factor must be between 1 and 4. gradient_accumulation_steps (`int`, *optional*, defaults to 1): Number of micro-batches to accumulate for the gradient calculation. Accumulates the gradient `gradient_accumulation` times before updating the model using the gradient. > Parameters related to parallelism layers_per_ipu (`List[int]`): Specifies the number of layers that will be put on each IPU for pipelined execution during training. For instance: `[2, 3, 4, 2]` specifies a 4-IPU pipeline, where the first two layers will be put on IPU0, the following three on IPU1, the next four on IPU2 and the last two on IPU3. If the default of [-1] is used, the layers will be split evenly over `ipus_per_replica` IPUs. The wildcard value '-1' can also be used in combination with integers. For instance: `[1, 2, -1, -1]` specifies a 4-IPU pipeline, where the first layer is put on IPU0, the next two layers on IPU1, and the remaining layers split evenly between IPU2 and IPU3. inference_layers_per_ipu (`List[int]`): Same as `layers_per_ipu` for inference only. ipus_per_replica (`int`, *optional*, defaults to `len(layers_per_ipu)`): Specifies the number of IPUs to use during training. This must be consistent with the number of IPUs used in `layers_per_ipu`. inference_ipus_per_replica (`int`, *optional*, defaults to `len(inference_layers_per_ipu) if ipus_per_replica==len(layers_per_ipu) else ipus_per_replica): Same as `ipus_per_replica` but for inference only. parallelize_kwargs (`Dict[str, Any]`, *optional*, defaults to None): Dictionary holding kwargs used for training model calls to `parallelize`. inference_parallelize_kwargs (`Dict[str, Any]`, *optional*, defaults to None): Dictionary holding kwargs used for inference model calls to `parallelize`. > Parameters for memory management optimizer_state_offchip (`bool`, *optional*, defaults to `True`): If `True`, uses the off-chip memory to store the optimizer state. If `False`, uses the on-chip memory. replicated_tensor_sharding (`bool`, *optional*, defaults to `False`): Shards the optimizer between replicas with zero-redundancy. matmul_proportion (`List[float]` or `float`, *optional*, defaults to 0.2): Sets the amount of temporary memory made available during training on per-IPU basis. Use this setting to control the amount of temporary memory available to operations such as: - convolution - matrix multiplication - embedding lookups - indexing operations inference_matmul_proportion (`List[float]` or `float`): Same as `matmul_proportion` for inference only. enable_half_partials (`bool`, *optional*, defaults to `True`): If `True`, sets the data type of partial results for matrix multiplication and convolution operators to float16. embedding_serialization_factor (`int`, *optional*, defaults to 1 if `serialized_embedding_splits_per_ipu` is `None`): The factor to use to serialize embeddings. Nothing happens if `embedding_serialization_factor = 1`. For `embedding_serialization_factor > 1`, the `torch.nn.Embedding` layer is replaced with a `optimum.graphcore.modeling_utils.SerializedEmbedding` layer. Note: only one of `embedding_serialization_factor` or `serialized_embedding_splits_per_ipu` should be provided. inference_embedding_serialization_factor (`int`, *optional*, defaults to 1 if `inference_serialized_embedding_splits_per_ipu` is `None`): Same as `embedding_serialization_factor` but for inference only. serialized_embedding_splits_per_ipu (`List[int]`, *optional*, defaults to None): Specifies the number of splits of the embedding layer that will be put on each IPU for pipelined execution. The format has to be the same as that for `layers_per_ipu` however wildcards are not supported. For instance: `[3, 1, 0, 0]` specifies how to place an embedding layer serialized into 4 sub-embedding layers across a 4-IPU pipeline. IPU-1 has 3 splits and IPU-2 has 1 split. The remaining IPUs have no sub-embedding layers. If an argument to this parameter is provided, it must: - be of the form `List[int>=0]` with atleast 1 split. - have the same pipeline length as `ipus_per_replica` - have splits that are consecutive with no zeros between splits e.g. `[3, 0, 2, 0]` is invalid - for generation, splits must lie entirely on the encoder or decoder portion of the pipeline. For example the 4-IPU pipeline `[3, 1, 0, 0]` for an encoder-decoder model can be split into `[3, 1]` and `[0, 0]`, however `[0, 1, 2, 0]` split into `[0, 1]` and `[2, 0]` is invalid. Note: only one of `embedding_serialization_factor` or `serialized_embedding_splits_per_ipu` should be set. inference_serialized_embedding_splits_per_ipu (`List[int]`, *optional*, defaults to None): Same as `serialized_embedding_splits_per_ipu` but for inference only. projection_serialization_factor (`int`, *optional*, defaults to 1 if `serialized_projection_splits_per_ipu` is `None`): The factor to use to either serialize the matmuls that are performed in the linear projection layer, or, serialize the projection layer into a set of individual linear layers that can be optionally placed on different IPUs. Nothing happens if `projection_serialization_factor = 1`. If `projection_serialization_factor > 1`, the `torch.nn.Linear` layer is replaced by a `optimum.graphcore.modeling_utils.SplitProjection` layer if `serialized_projection_splits_per_ipu` is provided and the linear layer's weights are not tied to another layer. Otherwise it is replaced by a `optimum.graphcore.modeling_utils.SerializedLinear` layer. Note: only one of `projection_serialization_factor` or `serialized_projection_splits_per_ipu` should be set. inference_projection_serialization_factor (`int`, *optional*, defaults to 1 if `inference_serialized_projection_splits_per_ipu` is `None`): Same as `projection_serialization_factor` but for inference only. serialized_projection_splits_per_ipu (`List[int]`, *optional*, defaults to None): Analogous to `serialized_embedding_splits_per_ipu`. Note: only one of `projection_serialization_factor` or `serialized_projection_splits_per_ipu` should be set. inference_serialized_projection_splits_per_ipu (`List[int]`, *optional*, defaults to None): Same as `serialized_projection_splits_per_ipu` but for inference only. recompute_checkpoint_every_layer (`bool`, *optional*, defaults to `False`): If `True`, uses gradient checkpointing at the end of every layer. It can help to reduce the memory impact. explicit_ir_inference (`bool`, *optional*, defaults to `False`): If `True`, uses experimental explicit-IR feature of PopART for inference models. This feature is only supported for inference models. For some cases explicit-IR can provide a better memory liveness schedule, reducing the peak memory during runtime. > Parameters related to host/device synchronization device_iterations (`int`, *optional*, defaults to 1): Number of iterations the device should run over the data before returning to the user during training. This is equivalent to running the IPU in a loop over the specified number of iterations, with a new batch of data each time. However, increasing the number of device iterations is more efficient because the loop runs on the IPU directly. inference_device_iterations (`int`, *optional*, defaults to 1): Same as `device_iterations` for inference. output_mode (`str`, *optional*, defaults to `"final"`): Specifies which data to return from a model. Allowed values: - `all`: returns a result for each batch. - `sum`: returns the sum of all batches. - `final`: returns the last batch. - `default`: `all` for inference, `final` for training. """ CONFIG_NAME = "ipu_config.json" FULL_CONFIGURATION_FILE = "ipu_config.json" class ManagedAttribute: def __init__(self, attr) -> None: self.attr = attr def __set__(self, obj, value): if isinstance(obj, IPUConfig): mode_attr = f"inference_{self.attr}" if obj.mode == "inference" else self.attr logger.debug(f"ManagedAttribute {self.attr} writing to {mode_attr}") return setattr(obj, mode_attr, value) def __get__(self, obj, objtype=None): if isinstance(obj, IPUConfig): mode_attr = f"inference_{self.attr}" if obj.mode == "inference" else self.attr logger.debug(f"ManagedAttribute {self.attr} reading from {mode_attr}") return getattr(obj, mode_attr) # Create descriptor based managed attributes which will either return the # `attribute` or `inference_attribute` versions of the attribute depending on the value of # `self.mode` ("training" by default). For example `_layers_per_ipu` # switches between `layers_per_ipu` and `inference_layers_per_ipu` modes = ("training", "inference") _layers_per_ipu = ManagedAttribute("layers_per_ipu") _ipus_per_replica = ManagedAttribute("ipus_per_replica") _matmul_proportion = ManagedAttribute("matmul_proportion") _embedding_serialization_factor = ManagedAttribute("embedding_serialization_factor") _serialized_embedding_splits_per_ipu = ManagedAttribute("serialized_embedding_splits_per_ipu") _projection_serialization_factor = ManagedAttribute("projection_serialization_factor") _serialized_projection_splits_per_ipu = ManagedAttribute("serialized_projection_splits_per_ipu") _parallelize_kwargs = ManagedAttribute("parallelize_kwargs") # Create a mapping of attributes to their list of validation functions attribute_validators = defaultdict(list) def _contents_geq_value_validator( name: str, value: Union[float, int, Sequence], floor_value: Union[float, int] ) -> None: """ Validates the values of Sequence and scalar types to be greater than `floor_value` For Sequence[Union[int, float]], ensure that all elements are >= floor_value For Union[float, int], ensure the scalar is >= floor_value """ # Do nothing for optional types if value is None: return elif isinstance(value, Sequence): if not all(elem >= floor_value for elem in value): raise ValueError( f"`IPUConfig` attribute `{name}` must have all elements >= {floor_value}. You provided {value=}" ) elif isinstance(value, (int, float)): if not value >= floor_value: raise ValueError(f"`IPUConfig` attribute `{name}` must be >= {floor_value}. You provided {value=}") else: raise ValueError( f"`contents_geq_value_validator` validates inputs of type:" f" Union[float, int, Sequence[Union[int, float]]]. You provided" f" attribute `{name}`, {value=}, {type(value)}" ) for attr, floor_value in ( ("layers_per_ipu", -1), ("inference_layers_per_ipu", -1), ("matmul_proportion", 0), ("inference_matmul_proportion", 0), ("replication_factor", 1), ("inference_replication_factor", 1), ("gradient_accumulation_steps", 1), ("ipus_per_replica", 1), ("inference_ipus_per_replica", 1), ("embedding_serialization_factor", 1), ("inference_embedding_serialization_factor", 1), ("projection_serialization_factor", 1), ("inference_projection_serialization_factor", 1), ("device_iterations", 1), ("inference_device_iterations", 1), ): attribute_validators[attr].append(partial(_contents_geq_value_validator, floor_value=floor_value)) def _output_mode_validator(name: str, value: str): allowed_values = ("all", "sum", "final", "default") if value not in allowed_values: raise ValueError( f"`IPUConfig` attribute `output_mode` can only take values in" f" {allowed_values}. You provided: {value=}" ) attribute_validators["output_mode"].append(_output_mode_validator) def _serialized_layer_splits_per_ipu_validator(name: str, value: int): """ Validates serialized_{projection/embedding}_splits_per_ipu attributes. If `value` is not None. `value` must be of type List[int>=0] with atleast 1 split on 1 IPU. Further splits in the pipeline must be consecutive. """ if value is None: return IPUConfig._contents_geq_value_validator(name, value, floor_value=0) # There must be atleast 1 split when the pipeline is provided if sum(value) < 1: raise ValueError(f"`IPUConfig` attribute `{name}={value}` must have atleast 1 split on 1 IPU.") # Check that splits are on consecutive IPUs (e.g. [3,0,2,0] is not allowed) for i, splits in enumerate(value[:-1]): if splits and value[i + 1] == 0 and sum(value[i + 1 :]) != 0: raise ValueError(f"`IPUConfig` attribute `{name}={value}` must have its splits on consecutive IPUs.") for attr in ( "serialized_embedding_splits_per_ipu", "inference_serialized_embedding_splits_per_ipu", "serialized_projection_splits_per_ipu", "inference_serialized_projection_splits_per_ipu", ): attribute_validators[attr].append(_serialized_layer_splits_per_ipu_validator) def __init__( self, replication_factor: int = 1, inference_replication_factor: int = 1, gradient_accumulation_steps: int = 1, layers_per_ipu: List[int] = [-1], inference_layers_per_ipu: Optional[List[int]] = None, ipus_per_replica: Optional[int] = None, inference_ipus_per_replica: Optional[int] = None, optimizer_state_offchip: bool = False, replicated_tensor_sharding: bool = False, matmul_proportion: Union[float, List[float]] = 0.2, inference_matmul_proportion: Optional[Union[float, List[float]]] = None, enable_half_partials: bool = True, embedding_serialization_factor: Optional[int] = None, inference_embedding_serialization_factor: Optional[int] = None, serialized_embedding_splits_per_ipu: Optional[List[int]] = None, inference_serialized_embedding_splits_per_ipu: Optional[List[int]] = None, projection_serialization_factor: Optional[int] = None, inference_projection_serialization_factor: Optional[int] = None, serialized_projection_splits_per_ipu: Optional[List[int]] = None, inference_serialized_projection_splits_per_ipu: Optional[List[int]] = None, recompute_checkpoint_every_layer: bool = False, device_iterations: int = 1, inference_device_iterations: int = 1, output_mode: str = "final", seed: Optional[int] = None, auto_loss_scaling: bool = False, executable_cache_dir: str = "", explicit_ir_inference: bool = False, parallelize_kwargs: Optional[Dict[str, Any]] = None, inference_parallelize_kwargs: Optional[Dict[str, Any]] = None, **kwargs, ): self.seed = seed # Default mode to `training` self.train() self.layers_per_ipu = layers_per_ipu self.inference_layers_per_ipu = inference_layers_per_ipu if inference_layers_per_ipu else self.layers_per_ipu self.ipus_per_replica = ipus_per_replica if ipus_per_replica else len(self.layers_per_ipu) # If ipus_per_replica is default, recalculate ipus_per_replica from inference_layers_per_ipu instead fallback_ipus_per_replica = self.ipus_per_replica if fallback_ipus_per_replica == len(self.layers_per_ipu) or self.inference_layers_per_ipu != [-1]: fallback_ipus_per_replica = len(self.inference_layers_per_ipu) self.inference_ipus_per_replica = ( inference_ipus_per_replica if inference_ipus_per_replica else fallback_ipus_per_replica ) self.matmul_proportion = matmul_proportion # If matmul_proportion is a list and its length is not equal to inference_ipus_per_replica, use the # default float value for matmul_proportion instead fallback_matmul_proportion = self.matmul_proportion if isinstance(self.matmul_proportion, list) and len(self.matmul_proportion) != self.inference_ipus_per_replica: fallback_matmul_proportion = 0.2 self.inference_matmul_proportion = ( inference_matmul_proportion if inference_matmul_proportion else fallback_matmul_proportion ) def check_and_set_replication_factor(attr_name, attr): if isinstance(attr, int): setattr(self, attr_name, attr) else: raise ValueError(f"{attr_name} must be of type `int`. You provided: {attr_name}={attr}, {type(attr)}.") check_and_set_replication_factor("replication_factor", replication_factor) check_and_set_replication_factor("inference_replication_factor", inference_replication_factor) # Non-transformer layers initialisation self.embedding_serialization_factor = embedding_serialization_factor self.inference_embedding_serialization_factor = inference_embedding_serialization_factor self.serialized_embedding_splits_per_ipu = serialized_embedding_splits_per_ipu self.inference_serialized_embedding_splits_per_ipu = inference_serialized_embedding_splits_per_ipu self.projection_serialization_factor = projection_serialization_factor self.inference_projection_serialization_factor = inference_projection_serialization_factor self.serialized_projection_splits_per_ipu = serialized_projection_splits_per_ipu self.inference_serialized_projection_splits_per_ipu = inference_serialized_projection_splits_per_ipu if kwargs.pop("sharded_execution_for_inference", None): warnings.warn( 'The "sharded_execution_for_inference" parameter is deprecated, sharded execution is always used during inference' ) if kwargs.pop("enable_half_first_order_momentum", None): warnings.warn('The "enable_half_first_order_momentum" parameter is deprecated') self.gradient_accumulation_steps = gradient_accumulation_steps self.device_iterations = device_iterations self.inference_device_iterations = inference_device_iterations self.optimizer_state_offchip = optimizer_state_offchip self.replicated_tensor_sharding = replicated_tensor_sharding self.auto_loss_scaling = auto_loss_scaling self.enable_half_partials = enable_half_partials self.executable_cache_dir = executable_cache_dir self.explicit_ir_inference = explicit_ir_inference self.embedding_serialization_factor = embedding_serialization_factor self.recompute_checkpoint_every_layer = recompute_checkpoint_every_layer self.output_mode = output_mode self.parallelize_kwargs = parallelize_kwargs or {} self.inference_parallelize_kwargs = inference_parallelize_kwargs or {} # TODO: remove this if unnecessary. self.execute_encoder_on_cpu_for_generation = kwargs.pop("execute_encoder_on_cpu_for_generation", False) # Raise error if user has provided unknown & unused kwarg if unknown_kwargs := (set(kwargs) - set(BaseConfig().to_dict())): raise IncompatibleIPUConfigError( "IPUConfig received unknown arguments:\n" + "\n".join([f" {k}={kwargs[k]}" for k in unknown_kwargs]) ) self._validate_ipu_config() @property def mode(self) -> str: return self._mode @mode.setter def mode(self, value): if value not in self.modes: raise ValueError( f"`IPUConfig` mode can only take values in {self.modes}." f" You provided: {value=}. Use the `train` and `eval` methods" " instead to avoid error." ) self._mode = value def train(self): self._mode = "training" return self def eval(self): self._mode = "inference" return self def _get_managed_attr_mode_name(self, attr: str) -> str: """ Returns the attribute name that a ManagedAttribute descriptor is currently referring to """ # Shallow check to ensure that the input attribute is actually # a managed attribute if hasattr(self, attr) and hasattr(self, f"inference_{attr}"): return attr if self.mode == "training" else f"inference_{attr}" # return attr if its not a managed attribute return attr def _get_attribute_type(self, name: str) -> Any: """ Returns the input `name` attribute type hints. Returns `Any` type by default. The return type for an attribute is only specific if it is a parameter in the signature of IPUConfig.__init__ """ try: type_hints = self._attribute_type_hints except AttributeError: type_hints = get_type_hints(IPUConfig.__init__) super().__setattr__("_attribute_type_hints", type_hints) return type_hints.get(name, Any) def __setattr__(self, name: str, value: Any): """Override __setattr__ to include value type checking and validation """ attr_type = self._get_attribute_type(name) try: typeguard.check_type(value, attr_type) except typeguard.TypeCheckError as e: raise TypeError( f"Setting `IPUConfig` attribute: {name}, type: {attr_type}" f" with {value=}, type: {type(value)} is invalid." ) from e # Run attribute value validators if name in self.attribute_validators: for vfunc in self.attribute_validators[name]: vfunc(name, value) return super().__setattr__(name, value) def _validate_ipu_config(self): """ Tests coherence of `IPUConfig` attributes for all modes in self.modes. For example if `matmul_proportion=[0.2, 0.2]`, `ipus_per_replica` must have value 2. Raises: IncompatibleIPUConfigError: Raised if any `IPUConfig` attributes are not coherent. """ if self.replicated_tensor_sharding and self.replication_factor == 1: logger.warning("`replicated_tensor_sharding` is not used when `replication_factor=1`") old_mode = self.mode for mode in self.modes: self.mode = mode ipus_per_replica_mode_str = self._get_managed_attr_mode_name("ipus_per_replica") # len(matmul_proportion) must equal ipus_per_replica if isinstance(self._matmul_proportion, list) and len(self._matmul_proportion) != self._ipus_per_replica: matmul_proportion_mode_str = self._get_managed_attr_mode_name("matmul_proportion") raise IncompatibleIPUConfigError( f"{matmul_proportion_mode_str}={self._matmul_proportion} should use the" f" same number of IPUs as {ipus_per_replica_mode_str}={self._ipus_per_replica}." ) # layers_per_ipu must have the same length as ipus per replica. # If wildcards are present in layers_per_ipu, let the call to `model.parallelize` # handle the validation if -1 not in self._layers_per_ipu and len(self._layers_per_ipu) != self._ipus_per_replica: layers_per_ipu_mode_str = self._get_managed_attr_mode_name("layers_per_ipu") raise IncompatibleIPUConfigError( f"{layers_per_ipu_mode_str}={self._layers_per_ipu} should use the" f" same number of IPUs as {ipus_per_replica_mode_str}={self._ipus_per_replica}." ) # Validate non-transformer layer placement configuration for layer in ("embedding", "projection"): mode_layer_splits_per_ipu_str = self._get_managed_attr_mode_name(f"serialized_{layer}_splits_per_ipu") mode_layer_splits_per_ipu = getattr(self, mode_layer_splits_per_ipu_str) mode_layer_serialisation_factor_str = self._get_managed_attr_mode_name(f"{layer}_serialization_factor") mode_layer_serialization_factor = getattr(self, mode_layer_serialisation_factor_str) # If the user has not provided either the layer_serialization_factor or # layer_splits_per_ipu, default the layer_serialization_factor to 1 if not (mode_layer_splits_per_ipu or mode_layer_serialization_factor): setattr(self, mode_layer_serialisation_factor_str, 1) # If the user provides both options, tell them only one is allowed and what each option is for if mode_layer_splits_per_ipu and mode_layer_serialization_factor: raise ValueError( f"Only one of `{mode_layer_serialisation_factor_str}` and `{mode_layer_splits_per_ipu_str}` should" f" be used at once. `{mode_layer_serialisation_factor_str}` should be used when you want your" f" {layer} layer serialised on the same IPU (which IPU depends on the model)." f" `{mode_layer_splits_per_ipu_str}` should be used when you want your {layer} layer to be split" " across multiple IPUs of your choice (or to choose which single IPU the layer is serialised on)." ) # Serialized layer splits per ipu pipeline must have the same pipeline length # as the number of ipus per replica if mode_layer_splits_per_ipu and len(mode_layer_splits_per_ipu) != self._ipus_per_replica: raise ValueError( f"{mode_layer_splits_per_ipu_str}={mode_layer_splits_per_ipu}" f" should use the same number of IPUs as {ipus_per_replica_mode_str}={self._ipus_per_replica}." ) self.mode = old_mode return self def _to_options(self, for_inference: bool = False, compile_only: bool = False) -> poptorch.Options: if not compile_only and poptorch.ipuHardwareVersion() not in (2, 21): raise RuntimeError("This requires an IPU Mk2 system to run.") if self.execute_encoder_on_cpu_for_generation: raise NotImplementedError("execute_encoder_on_cpu_for_generation is not supported yet.") old_mode = self.mode self.eval() if for_inference else self.train() opts = Options() opts.autoRoundNumIPUs(True) opts.replicationFactor(self.inference_replication_factor if for_inference else self.replication_factor) opts.deviceIterations(self.inference_device_iterations if for_inference else self.device_iterations) if not for_inference: # Set gradient accumulation factor opts.Training.gradientAccumulation(self.gradient_accumulation_steps) opts.Training.accumulationAndReplicationReductionType(poptorch.ReductionType.Mean) # Enable automatic loss scaling # Note that this is an experimental feature. Note also that it expects # accumulationAndReplicationReductionType to be set to Mean as above, # and for accumulation by the optimizer to be done in half precision # using accum_type=torch.float16 during optimizer instantiation. if self.auto_loss_scaling and not for_inference: opts.Training.setAutomaticLossScaling(True) # Return all results from IPU to host output_mode_mapping = { "all": OutputMode.All, "sum": OutputMode.Sum, "final": OutputMode.Final, "default": OutputMode.Default, } training_output_mode = output_mode_mapping.get(self.output_mode, None) if training_output_mode is None: supported_output_modes = ", ".join(output_mode_mapping.keys()) raise KeyError( f"{self.output_mode} is not a valid poptorch.OutputMode, supported output modes: {supported_output_modes}" ) opts.outputMode(OutputMode.All if for_inference else training_output_mode) if self.seed: opts.randomSeed(self.seed) # Enable replicated tensor sharding of optimizer state # with optimizer state residing either on-chip or in DRAM. # RTS is only enabled if replication factor is also greater than 1 opts.TensorLocations.setOptimizerLocation( poptorch.TensorLocationSettings() # Optimizer state lives on- or off-chip .useOnChipStorage(not self.optimizer_state_offchip) # Shard optimizer state between replicas with zero-redundancy .useReplicatedTensorSharding(self.replicated_tensor_sharding and opts.replication_factor > 1) ) if for_inference: opts.setExecutionStrategy(poptorch.ShardedExecution(poptorch.AutoStage.AutoIncrement)) else: # Use Pipelined Execution opts.setExecutionStrategy(poptorch.PipelinedExecution(poptorch.AutoStage.AutoIncrement)) # Compile offline (no IPUs required) if compile_only: opts.useOfflineIpuTarget() matmul_proportion = copy.deepcopy(self._matmul_proportion) if isinstance(matmul_proportion, float): matmul_proportion = [matmul_proportion] * self._ipus_per_replica mem_prop = {f"IPU{i}": matmul_proportion[i] for i in range(self._ipus_per_replica)} opts.setAvailableMemoryProportion(mem_prop) # Enable caching the compiled executable to disk if self.executable_cache_dir and self.executable_cache_dir != "disabled": opts.enableExecutableCaching(self.executable_cache_dir) opts._Popart.set("saveInitializersToFile", NamedTemporaryFile().name) # Enable stochastic rounding (recommended for training with FP16) opts.Precision.enableStochasticRounding(not for_inference) # Half precision partials for matmuls and convolutions if self.enable_half_partials: opts.Precision.setPartialsType(torch.float16) # PopART performance options # # Only stream needed tensors back to host opts._Popart.set("disableGradAccumulationTensorStreams", True) # Parallelize optimizer step update across IPUs opts._Popart.set( "accumulateOuterFragmentSettings.schedule", int(popart.AccumulateOuterFragmentSchedule.OverlapMemoryOptimized), ) opts._Popart.set("accumulateOuterFragmentSettings.excludedVirtualGraphs", ["0"]) # Enable patterns for better throughput and memory reduction opts._Popart.set("outlineThreshold", 10.0) opts._Popart.set("subgraphCopyingStrategy", int(popart.SubgraphCopyingStrategy.JustInTime)) opts._Popart.set("scheduleNonWeightUpdateGradientConsumersEarly", True) opts._Popart.setPatterns( {"TiedGather": True, "TiedGatherAccumulate": True, "UpdateInplacePrioritiesForIpu": True} ) # Options for profiling with Popvision engine_options = { "opt.useAutoloader": "true", "target.syncReplicasIndependently": "true", } if for_inference and self.explicit_ir_inference: opts._popart.set("enableExplicitIR", True) opts._Popart.set("engineOptions", engine_options) self.mode = old_mode return opts def to_options(self, for_inference: bool = False, compile_only: bool = False) -> poptorch.Options: """ Creates a `poptorch.Options` instance from the `IPUConfig` instance. Args: for_inference (`bool`, defaults to `False`): If `True`, the resulting `poptorch.Options` will be adapted for inference. If `False`, the resulting `poptorch.Options` will be adapted for training. compile_only (`bool`, defaults to `False`): If True, compilation will be performed offline, no IPUs required. Returns: `poptorch.Options`: The options representing the `IPUConfig` instance. """ return self._to_options(for_inference=for_inference, compile_only=compile_only) # Adapted from BaseConfig.to_dict def to_dict(self) -> Dict[str, Any]: """ Serializes this instance to a Python dictionary. Returns: `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. """ output = super().to_dict() # Remove type hints as they are not serializable output.pop("_attribute_type_hints", None) # Remove mode as it's not relevant for a dict output.pop("_mode", None) return output def batch_size_factor(self, for_inference: bool = False) -> int: """ Computes the factor to apply to the micro batch size to calculate the combined batch size. Args: for_inference (`bool`, defaults to `False`): Whether the factor is being use to compute the batch size for inference or not. Returns: `int`: The batch size factor. """ replication_factor = self.inference_replication_factor if for_inference else self.replication_factor gradient_accumulation_steps = 1 if for_inference else self.gradient_accumulation_steps device_iterations = self.inference_device_iterations if for_inference else self.device_iterations return replication_factor * gradient_accumulation_steps * device_iterations def update_from_string(self, update_str: str): """ Updates attributes of the `IPUConfig` class with attributes from `update_str`. The expected format is ints, floats and strings as is, and for booleans use `true` or `false`, and for lists use `[a b c d]`. For example: `"n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index, matmul_proportion=[0.08 0.2 0.25 0.25]"`. The keys to change must already exist in the config object. Args: update_str (`str`): String with attributes that should be updated for this class. """ d = dict(x.split("=") for x in update_str.split(",")) for k, v in d.items(): if not hasattr(self, k): raise ValueError(f"Key {k} isn't in the original config dict") old_v = getattr(self, k) if isinstance(old_v, bool): if v.lower() in ["true", "1", "y", "yes"]: v = True elif v.lower() in ["false", "0", "n", "no"]: v = False else: raise ValueError(f"Can't derive true or false from {v} (key {k})") elif isinstance(old_v, int): v = int(v) elif isinstance(old_v, float): v = float(v) elif isinstance(old_v, list): v = json.loads(v.replace(" ", ",")) elif not isinstance(old_v, str): raise ValueError( f"You can only update int, float, bool, list or string values in the config, got {v} for key {k}" ) setattr(self, k, v)