squad_training.py

# Copyright (c) 2021 Graphcore Ltd. All rights reserved.

import logging
import time
import numpy as np
from typing import Dict, Union

import popxl
from popxl import ops
import popxl_addons as addons
from popxl_addons import TaskSession
from popxl_addons.optimizers.adam import AdamOptimizerStep
from popxl_addons.graph import GraphWithNamedArgs
from popxl_addons.variable_factory import NamedVariableFactories
from popxl_addons.named_tensors import NamedTensors
from popxl_addons.transforms.batch_serialisation import (
    batch_serialise_fwd_and_grad,
    batch_serial_buffer,
    batch_serialise,
)
from popxl_addons.transforms.repeat_graph import repeat_graph
from popxl_addons.rts import (
    all_gather_replica_sharded_graph,
    reduce_replica_sharded_graph,
    replica_sharded_spec,
)

from popxl_addons.remote import (
    named_variable_buffers,
    load_remote_graph,
    store_remote_graph,
)
from popxl_addons.named_replica_grouping import NamedReplicaGrouping, get_ild_replica_grouping

from config import BertConfig, CONFIG_DIR
from utils.setup import bert_config_setup
from modelling.embedding import BertEmbeddings
from modelling.bert_model import BertLayer
from modelling.squad import BertSquadLossAndGrad

__all__ = ["squad_training_phased"]

OptimGraphs = Dict[str, GraphWithNamedArgs]
RTS_THRESHOLD = 1024


def requires_weight_decay(t: popxl.Tensor):
    return not any(map(lambda exclude: exclude in t.name, ["norm", "bias"]))


def get_rts_groups(facts: NamedVariableFactories) -> NamedReplicaGrouping:
    ir = popxl.gcg().ir

    rts_groups = {}
    for k, f in facts.to_dict().items():
        size = np.prod(f.shape)
        rg = f.replica_grouping
        # Limit RTS to within an ILD
        rg = get_ild_replica_grouping(rg)
        if size % rg.group_size == 0 and size >= RTS_THRESHOLD:
            rts_groups[k] = rg
        else:
            rts_groups[k] = ir.replica_grouping(group_size=1)
    return NamedReplicaGrouping.from_dict(rts_groups)


class Graphs:
    def __init__(self, config: BertConfig, layer: addons.Module, optimizer: addons.Module, *args, **kwargs):
        # Create Graphs for computing forward, gradient and optimizer
        fwd_args, self.fwd = layer.create_graph(*args, **kwargs)
        required_grads = () if isinstance(layer, BertEmbeddings) else (self.fwd.graph.inputs[0],)
        grad_args, self.grad = addons.autodiff_with_accumulation(
            self.fwd, self.fwd.args.tensors, grads_required=required_grads
        )

        optim_args = {}
        self.optim: OptimGraphs = {}
        fwd_rts_groups = get_rts_groups(fwd_args)
        shard_over = fwd_rts_groups.to_dict()
        for name, var in self.fwd.args.to_dict().items():
            optim_args[name], self.optim[name] = optimizer.create_graph(
                replica_sharded_spec(var, shard_over[name]),
                replica_sharded_spec(var, shard_over[name]),
                lr=popxl.TensorSpec((), popxl.float32),
                weight_decay=config.training.optimizer.weight_decay if requires_weight_decay(var) else 0.0,
                beta1=config.training.optimizer.beta1,
                beta2=config.training.optimizer.beta2,
                eps=(1e-6 * config.execution.loss_scaling),
                bias_correction=False,
            )

        optim_args = NamedVariableFactories.from_dict(optim_args)
        # Variables required
        self.args = NamedVariableFactories(fwd=fwd_args, optim=optim_args)
        self.grad_args = grad_args

        # Create remote buffers for fwd vars and optimiser state.
        entries = config.model.layers if isinstance(layer, BertLayer) else 1
        rts_group = NamedReplicaGrouping(fwd=fwd_rts_groups, optim=get_rts_groups(optim_args))
        shard_over = {k: rg.group_size for k, rg in rts_group.to_dict().items()}
        self.buffers = named_variable_buffers(self.args, entries, shard_over_dict=shard_over)

        # Create Graphs for loading/gathering/storing/reducing
        self._fwd_load, self._fwd_load_names = load_remote_graph(self.buffers.fwd, entries)
        self._optim_load, self._optim_load_names = load_remote_graph(self.buffers, entries)
        self._optim_store = store_remote_graph(self.buffers, entries)

        self._fwd_all_gather, self._fwd_all_gather_names = all_gather_replica_sharded_graph(
            NamedTensors.pack(self._fwd_load_names, self._fwd_load.graph.outputs)
        )

        grad_accums = self.grad.args.copy()
        grad_accums.pop("mean_accum_counter")
        self._grad_reduce, self._grad_reduce_names = reduce_replica_sharded_graph(
            grad_accums, "mean", shard_groups=get_rts_groups(self.grad_args)
        )

    @classmethod
    def empty(cls):
        return super().__new__(cls)

    def fwd_load(self, i: Union[int, popxl.Tensor]):
        return NamedTensors.pack(self._fwd_load_names, self._fwd_load.call(i))

    def optim_load(self, i: Union[int, popxl.Tensor]):
        return NamedTensors.pack(self._optim_load_names, self._optim_load.call(i))

    def optim_store(self, args: NamedTensors, i: Union[int, popxl.Tensor]):
        return self._optim_store.bind(args).call(i)

    def fwd_all_gather(self, args: NamedTensors):
        return NamedTensors.pack(self._fwd_all_gather_names, self._fwd_all_gather.bind(args).call())

    def grad_reduce(self, args: NamedTensors):
        return NamedTensors.pack(self._grad_reduce_names, self._grad_reduce.bind(args).call())


def create_squad_graph(config: BertConfig, optimizer: addons.Module, *args, **kwargs):
    """Squad combines the forward, loss and grad into a single Module."""
    args, graph = BertSquadLossAndGrad(config).create_graph(*args, **kwargs)

    optim_args: Dict[str, NamedVariableFactories] = {}
    optim_graphs: OptimGraphs = {}
    shard_over = get_rts_groups(args.fwd).to_dict()
    for name, var in graph.args.fwd.to_dict().items():
        optim_args[name], optim_graphs[name] = optimizer.create_graph(
            replica_sharded_spec(var, shard_over[name]),
            replica_sharded_spec(var, shard_over[name]),
            lr=popxl.constant(1e-3),  # TODO: Replace with TensorSpec
            weight_decay=config.training.optimizer.weight_decay if requires_weight_decay(var) else 0.0,
            beta1=config.training.optimizer.beta1,
            beta2=config.training.optimizer.beta2,
            eps=(1e-6 * config.execution.loss_scaling),
            bias_correction=False,
        )

    args.insert("optim", NamedVariableFactories.from_dict(optim_args))
    grad_args = args.pop("grad")

    squad = Graphs.empty()
    squad.fwd = graph
    squad.optim = optim_graphs
    squad.args = args
    squad.grad_args = grad_args
    # Create remote buffers for fwd vars and optimiser state.
    rts_group = get_rts_groups(args)
    shard_over = {k: rg.group_size for k, rg in rts_group.to_dict().items()}
    squad.buffers = named_variable_buffers(args, shard_over_dict=shard_over)
    squad._fwd_load, squad._fwd_load_names = load_remote_graph(squad.buffers.fwd)
    squad._optim_load, squad._optim_load_names = load_remote_graph(squad.buffers)
    squad._optim_store = store_remote_graph(squad.buffers)
    squad._fwd_all_gather, squad._fwd_all_gather_names = all_gather_replica_sharded_graph(
        NamedTensors.pack(squad._fwd_load_names, squad._fwd_load.graph.outputs)
    )
    grad_accums = grad_args.copy()
    grad_accums.pop("mean_accum_counter")
    squad._grad_reduce, squad._grad_reduce_names = reduce_replica_sharded_graph(
        grad_accums, "mean", shard_groups=get_rts_groups(grad_accums)
    )

    return squad


def get_optimizer_state(name: str, state: NamedTensors) -> NamedTensors:
    attrs = name.split(".")
    for attr in attrs:
        state = getattr(state, attr)
    return state


def optimizer_step(optim_graphs: OptimGraphs, vars_and_state: NamedTensors, grads: NamedTensors, lr: popxl.Tensor):
    _variables = vars_and_state.fwd.to_dict()
    _state = vars_and_state.optim
    _grads = grads.accum.to_dict()
    for name, graph in optim_graphs.items():
        graph.bind(get_optimizer_state(name, _state)).call(_variables[name], _grads[name], lr)


def embeddings_batch_serialise(
    config: BertConfig,
    embeddings: Graphs,
    inputs: addons.InputStreams,
    x_buffer: popxl.RemoteBuffer,
    dx_buffer: popxl.RemoteBuffer,
):
    fwd, grad = batch_serialise_fwd_and_grad(
        embeddings.fwd,
        embeddings.grad,
        embeddings.fwd.args,
        config.gradient_accumulation,
        load_handles={
            embeddings.fwd.graph.inputs[0]: inputs.words,
            embeddings.fwd.graph.inputs[1]: inputs.token_type,
            embeddings.grad.graph.inputs[0]: (dx_buffer, 0),
        },
        store_streams={},
        store_buffers={embeddings.fwd.graph.outputs[0]: (x_buffer, 0)},
        seed_input=embeddings.fwd.graph.inputs[2],
        rows=1,
        io_mode="io",
    )
    embeddings.fwd = fwd.graph
    embeddings.grad = grad.graph


def layer_batch_serialise(
    config: BertConfig,
    layer: Graphs,
    x_buffer: popxl.RemoteBuffer,
    dx_buffer: popxl.RemoteBuffer,
    mask_buffer: popxl.RemoteBuffer,
):
    fwd, grad = batch_serialise_fwd_and_grad(
        layer.fwd,
        layer.grad,
        layer.fwd.args,
        config.gradient_accumulation,
        load_handles={
            layer.fwd.graph.inputs[0]: (x_buffer, 0),
            layer.fwd.graph.inputs[1]: (mask_buffer, None),
            layer.grad.graph.inputs[0]: (dx_buffer, 1),
        },
        store_streams={},
        store_buffers={layer.fwd.graph.outputs[0]: (x_buffer, 1), layer.grad.graph.outputs[0]: (dx_buffer, 0)},
        seed_input=layer.fwd.graph.inputs[2],
        rows=config.model.layers,
        io_mode="io",
    )
    layer.fwd = fwd.graph
    layer.grad = grad.graph


def squad_batch_serialise(
    config: BertConfig,
    squad: Graphs,
    inputs: addons.InputStreams,
    outputs: addons.OutputStreams,
    x_buffer: popxl.RemoteBuffer,
    dx_buffer: popxl.RemoteBuffer,
):
    bs_squad = batch_serialise(
        squad.fwd,
        config.gradient_accumulation,
        load_handles={
            squad.fwd.graph.inputs[0]: (x_buffer, config.model.layers),
            squad.fwd.graph.inputs[1]: inputs.labels,
        },
        store_streams={squad.fwd.graph.outputs[0]: outputs.loss},
        store_buffers={squad.fwd.graph.outputs[1]: (dx_buffer, config.model.layers)},
        rows=1,
        io_mode="io",
    )
    squad.fwd = bs_squad.graph


def squad_training_phased(config: BertConfig) -> TaskSession:
    assert config.execution.data_parallel > 1

    ir = popxl.Ir()
    ir.replication_factor = config.execution.data_parallel
    opts = ir._pb_ir.getSessionOptions()
    opts.numIOTiles = config.execution.io_tiles
    opts.enableStochasticRounding = config.training.stochastic_rounding
    opts.partialsTypeMatMuls = "half"

    t = time.time()
    main = ir.main_graph

    with main:
        # -----  Define input and output streams -----
        input_shape = (config.execution.micro_batch_size * config.model.sequence_length,)
        inputs = addons.InputStreams(
            words=(input_shape, popxl.uint32),
            token_type=(input_shape, popxl.uint32),
            mask=(input_shape, config.model.dtype),
            labels=((config.execution.micro_batch_size, 2), popxl.uint32),
            lr=((), popxl.float32),
        )
        outputs = addons.OutputStreams(loss=((), config.model.dtype))

        # ---- Initialise Random Seed ----
        seed_v, seed = addons.seed_variable(config.model.seed)

        # ----- Build compute graphs -----
        optimizer = AdamOptimizerStep(cache=True)

        embeddings = Graphs(
            config, BertEmbeddings(config), optimizer, inputs.words.spec, inputs.token_type.spec, seed=seed.spec
        )
        layer = Graphs(
            config, BertLayer(config), optimizer, embeddings.fwd.graph.outputs[0].spec, inputs.mask.spec, seed=seed.spec
        )

        squad = create_squad_graph(config, optimizer, layer.fwd.graph.outputs[0].spec, inputs.labels.spec)

        # ---- Transform graphs ----

        # Recomputation
        embeddings.grad = addons.recompute_graph(embeddings.grad)
        layer.grad = addons.recompute_graph(layer.grad)

        # Batch Serialisation
        #   Buffers
        x_buffer = batch_serial_buffer(
            embeddings.fwd.graph.outputs[0], steps=config.gradient_accumulation, rows=config.model.layers + 1
        )
        dx_buffer = batch_serial_buffer(
            embeddings.grad.graph.inputs[0], steps=config.gradient_accumulation, rows=config.model.layers + 1
        )
        mask_buffer = batch_serial_buffer(layer.fwd.graph.inputs[1], steps=config.gradient_accumulation)
        #   Graphs
        embeddings_batch_serialise(config, embeddings, inputs, x_buffer, dx_buffer)
        layer_batch_serialise(config, layer, x_buffer, dx_buffer, mask_buffer)
        squad_batch_serialise(config, squad, inputs, outputs, x_buffer, dx_buffer)

        # Available Memory Proportion
        addons.set_available_memory_proportion_by_ipu(ir, config.execution.available_memory_proportion)

        # ----- Create Variables -----

        variables = NamedTensors(random_seed=seed_v)
        variables.insert("embeddings", embeddings.args.init_remote(embeddings.buffers, 0, "embeddings"))
        variables.insert("squad", squad.args.init_remote(squad.buffers, 0, "squad"))
        variables.insert(
            "layer",
            NamedTensors.from_dict(
                {n: layer.args.init_remote(layer.buffers, n, f"layer.{n}") for n in range(config.model.layers)}
            ),
        )

        # ---- Execute ----

        with popxl.in_sequence():
            # Load current learning rate
            lr = ops.host_load(inputs.lr)
            # Increment random seed
            seed += 1

            @popxl.io_tiles()
            def fill_buffer_from_host(i: popxl.Tensor, stream: popxl.HostToDeviceStream, buffer: popxl.RemoteBuffer):
                t = ops.host_load(stream)
                ops.remote_store(buffer, i, t)

            # Load from host then store all masks TODO: Move this into the embedding loop
            mask_fill_graph = ir.create_graph(
                fill_buffer_from_host, popxl.constant(0, popxl.uint32), inputs.mask, mask_buffer
            )
            for i in range(config.gradient_accumulation):
                ops.call(mask_fill_graph, i)

            def embedding_fwd_phase(seed):
                # Load Embedding layer
                embeddings_vars = embeddings.fwd_load(0)
                embeddings_vars = embeddings.fwd_all_gather(embeddings_vars)
                # Forward
                seed, embed_seed = ops.split_random_seed(seed)
                embeddings.fwd.bind(embeddings_vars).call(0, embed_seed)
                return seed

            seed = embedding_fwd_phase(seed)

            def single_bert_layer_fwd_phase(n: popxl.Tensor, seed: popxl.Tensor):
                # Load Encoder layers
                layer_vars = layer.fwd_load(n)
                layer_vars = layer.fwd_all_gather(layer_vars)
                # Forward
                seed, layer_seed = ops.split_random_seed(seed)
                layer.fwd.bind(layer_vars).call(n, layer_seed)
                return n + 1, seed

            i = popxl.constant(0)
            bwd_graph = ir.create_graph(single_bert_layer_fwd_phase, i, seed)
            ops.repeat(bwd_graph, config.model.layers, i, seed)

            def squad_fwd_grad_optimizer_phase():
                # Load Squad layer
                squad_vars = squad.optim_load(0)
                squad_fwd_vars = NamedTensors(
                    fwd=squad.fwd_all_gather(squad_vars.fwd), grad=squad.grad_args.init_zero()
                )
                # Forward + Gradient
                squad.fwd.bind(squad_fwd_vars).call(0)
                # Optimizer
                reduced_grads = squad.grad_reduce(squad_fwd_vars.grad)
                optimizer_step(squad.optim, squad_vars, reduced_grads, lr)
                # Store
                squad.optim_store(squad_vars, 0)

            squad_fwd_grad_optimizer_phase()

            def single_bert_layer_grad_optimizer_phase(n: popxl.Tensor, lr: popxl.Tensor):
                layer_vars = layer.optim_load(n)
                layer_fwd_vars = layer.fwd_all_gather(layer_vars.fwd)
                # Gradient
                grads = layer.grad_args.init_zero()
                bwd_vars = grads.copy()
                bwd_vars.update(layer_fwd_vars)
                layer.grad.bind(bwd_vars).call(n)
                # Optimizer
                reduced_grads = layer.grad_reduce(grads)
                optimizer_step(layer.optim, layer_vars, reduced_grads, lr)
                # Store
                layer.optim_store(layer_vars, n)
                return n - 1

            i = popxl.constant(config.model.layers - 1)
            bwd_graph = ir.create_graph(single_bert_layer_grad_optimizer_phase, i, lr)
            ops.repeat(bwd_graph, config.model.layers, i, lr)

            def embedding_grad_optimizer_phase():
                # Load Embeddings layer
                embeddings_vars = embeddings.optim_load(0)
                embeddings_fwd_vars = embeddings.fwd_all_gather(embeddings_vars.fwd)
                # Gradient
                grads = embeddings.grad_args.init_zero()
                bwd_vars = grads.copy()
                bwd_vars.update(embeddings_fwd_vars)
                embeddings.grad.bind(bwd_vars).call(0)
                # Optimizer
                reduced_grads = embeddings.grad_reduce(grads)
                optimizer_step(embeddings.optim, embeddings_vars, reduced_grads, lr)
                # Store
                embeddings.optim_store(embeddings_vars, 0)

            embedding_grad_optimizer_phase()

    repeat_graph(main, config.execution.device_iterations)

    fwd_vars = NamedTensors(
        embeddings=variables.embeddings.fwd,
        layer=NamedTensors.from_dict({i: variables.layer[i].fwd for i in range(config.model.layers)}),
        squad=variables.squad.fwd,
    )

    logging.info(f"popxl IR construction duration: {(time.time() - t) / 60:.2f} mins")

    ir.num_host_transfers = config.execution.device_iterations * config.gradient_accumulation
    session = TaskSession(inputs=inputs, outputs=outputs, state=fwd_vars, ir=ir, device_desc="ipu_hw")
    return session


def main():
    """Run a benchmark configuration"""
    config, _ = bert_config_setup(CONFIG_DIR / "squad_training.yml", "phased", "large")
    session = squad_training_phased(config)

    inputs = {
        stream: np.ones(session._full_input_shape(stream.shape), stream.dtype.as_numpy())
        for stream in session.expected_inputs()
    }

    with session:
        # Skip one result
        session.run(inputs)

        durations = []
        for _ in range(5):
            start = time.time()
            session.run(inputs)
            durations.append(time.time() - start)
    duration = np.mean(durations)

    samples_per_step = config.execution.device_iterations * config.training.global_batch_size
    result_str = f"Duration: {duration} s " f"throughput: {samples_per_step/duration:6.1f} samples/sec "
    logging.info(result_str)


if __name__ == "__main__":
    main()