# Copyright (c) 2023 Graphcore Ltd. All rights reserved. import logging import time import numpy as np from typing import Dict, Union, Type, Optional import popdist import popxl from popxl import ops, gcg import popxl_addons as addons from pretraining_config import ( RTS_THRESHOLD, RTS_ACTIVATIONS_THRESHOLD, USE_IO_TILES, GraphConf, PhaseConf, gen_layer_config, filter, RTS_ACT, ) from popxl_addons.optimizers.adam import AdamOptimizerStep from popxl_addons import TaskSession from popxl_addons.utils import OrderedDict, timer from popxl_addons.patterns import apply_pre_alias_patterns from popxl_addons.graph import GraphWithNamedArgs from popxl_addons.variable_factory import NamedVariableFactories from popxl_addons.named_replica_grouping import NamedReplicaGrouping from popxl_addons.named_tensors import NamedTensors from popxl_addons.transforms.repeat_graph import repeat_graph from popxl_addons.transforms.batch_serialisation import ( batch_serialise_fwd_and_grad, batch_serial_buffer, batch_serialise, RemoteHandle, ) from popxl_addons.rts import ( all_gather_replica_sharded_graph, replica_sharded_spec, reduce_replica_sharded_graph, reduce_replica_sharded_tensor, ) from popxl_addons.remote import ( named_variable_buffers, load_remote_graph, store_remote_graph, create_remote_buffer, NamedRemoteBuffers, ) from popxl_addons.ops.grad_reduce_square_add import grad_reduce_square_add from config import GPTConfig, CONFIG_DIR from utils.setup import gpt_config_setup from modelling.embedding import GPTEmbeddingsTP, generate_positions from modelling.decoder import GPTDecoderBlockTP from modelling.gpt_lm import GPTLMHeadLossTP, HeadFwdBwdTiedEmb from utils.utils import replica_groups __all__ = ["pretraining"] OptimGraphs = Dict[str, GraphWithNamedArgs] def get_activ_shard_group(a: popxl.Tensor, shard_group: popxl.ReplicaGrouping, RTS_ACTIVATIONS_THRESHOLD: int): return shard_group if a.nelms >= RTS_ACTIVATIONS_THRESHOLD else popxl.gcg().ir.replica_grouping(group_size=1) 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.const_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) def requires_weight_decay(t: popxl.Tensor): return not any(map(lambda exclude: exclude in t.name, ["norm", "bias"])) def optimizer_graphs( config: GPTConfig, optimizer: addons.Module, variables: NamedTensors, facts: NamedVariableFactories ): optim_facts = {} optim_graphs = {} replica_groups = facts.replica_groupings.to_dict() rts_groups = get_rts_groups(facts) for name, var in variables.to_dict().items(): input_spec = replica_sharded_spec(var, shard_over=rts_groups[name]) replica_group = replica_groups[name].const_rg optim_facts[name], optim_graphs[name] = optimizer.create_graph( input_spec, input_spec, lr=popxl.TensorSpec((), popxl.float32), replica_grouping=replica_group, 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, bias_correction=True, first_order_dtype=popxl.float32, loss_scaling=config.execution.loss_scaling, global_norm=popxl.TensorSpec((), popxl.float32), global_norm_max=config.training.optimizer.gradient_clipping, ) return NamedVariableFactories.from_dict(optim_facts), optim_graphs class Graphs: def __init__( self, config: GPTConfig, layer_configs, optimizer: addons.Module, Layer: Type[addons.Module], entries: int, reuse_buffers: Optional[Dict] = None, *args, **kwargs, ): self.config = config self.Layer = Layer _, rg_dp = replica_groups(config) self.layer_config = layer_configs[Layer] graph_settings: GraphConf = self.layer_config.graph_config # Create Graphs for computing forward, gradient and optimizer fwd_facts, self.fwd = Layer(config).create_graph(*args, **kwargs) # Autodiff # self.fwd.args only include named tensors/variables tensors_to_accum = filter(self.fwd.args, graph_settings.accumulate) grads_provided = filter(self.fwd.graph.outputs, graph_settings.grads_provided) grads_required = filter(self.fwd.graph.inputs, graph_settings.grads_required) called_graphs_grad_info = {} if config.execution.attention_serialisation > 1 and Layer == GPTDecoderBlockTP: # Optimisation to recompute each blk separately assert len(self.fwd.graph.called_graphs) == 1, "expected exactly 1 called graph by decoder layer fwd" blk_graph = GraphWithNamedArgs(self.fwd.graph.called_graphs[0]) grad_blk_graph = addons.transforms.autodiff(blk_graph, grads_required=blk_graph.graph.inputs[:-2]) grad_blk_graph = addons.transforms.recompute_graph(grad_blk_graph) called_graphs_grad_info[blk_graph.graph] = grad_blk_graph.grad_graph_info grad_facts, self.bwd = addons.autodiff_with_accumulation( self.fwd, tensors_to_accum.values_flat(), grads_provided=grads_provided, grads_required=grads_required, replica_groupings=fwd_facts.replica_groupings, called_graphs_grad_info=called_graphs_grad_info, ) popxl.transforms.decompose_sum(self.bwd.graph) reuse_rg = {} if graph_settings.reuse: assert len(graph_settings.reuse) == len(reuse_buffers) for var_name in graph_settings.reuse: assert var_name in reuse_buffers, f"{var_name} not in {reuse_buffers}" grad_facts.accum.pop(var_name) tensors_to_accum.pop(var_name) fwd_fact = fwd_facts.pop(var_name) reuse_rg[var_name] = fwd_fact.replica_grouping # Optimiser optim_facts, self.optim = optimizer_graphs(config, optimizer, tensors_to_accum, fwd_facts) # Variables required self.facts = NamedVariableFactories(fwd=fwd_facts, optim=optim_facts) self.grad_facts = grad_facts remote_buffer_facts = NamedVariableFactories() if graph_settings.remote_buffer_fwd: remote_buffer_facts.insert("fwd", fwd_facts.copy()) if graph_settings.remote_buffer_bwd: remote_buffer_facts.insert("bwd", grad_facts.copy()) remote_buffer_facts.bwd.pop("mean_accum_counter") if graph_settings.remote_buffer_optim: remote_buffer_facts.insert("optim", optim_facts.copy()) rts_groups = get_rts_groups(remote_buffer_facts) shard_over = {k: rg.group_size for k, rg in rts_groups.to_dict().items()} self.buffers = named_variable_buffers(remote_buffer_facts, shard_over_dict=shard_over) self.remote_buffer_facts = remote_buffer_facts ### Create Graphs for loading/gathering/storing/reducing remote buffers # Store fwd and optim self._optim_fwd_store = store_remote_graph(self.buffers.filter_keys(["fwd", "optim"]), entries) # Store bwd if "bwd" in self.buffers: self._grad_store = store_remote_graph(self.buffers.bwd, entries) # Load fwd if "fwd" in self.buffers: fwd_buffers: NamedRemoteBuffers = self.buffers.fwd.copy() if graph_settings.reuse: for var_name in graph_settings.reuse: fwd_buffers.insert(var_name, reuse_buffers[var_name], overwrite=True) rts_groups.fwd.insert(var_name, reuse_rg[var_name], overwrite=True) self._fwd_load, self._fwd_load_names = load_remote_graph(fwd_buffers, entries) # Load optim + fwd self._optim_fwd_load, self._optim_fwd_load_names = load_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), replica_groups=rts_groups.fwd, use_io_tiles=USE_IO_TILES, ) # RTS graph: reduce grad_accums = self.bwd.args.copy() grad_accums.pop("mean_accum_counter") if graph_settings.reuse: for var_name in graph_settings.reuse: grad_accums.accum.pop(var_name) rts_bwd_group = NamedReplicaGrouping(accum=rts_groups.fwd.copy()) self._grad_reduce, self._grad_reduce_names = reduce_replica_sharded_graph( grad_accums, "mean", shard_groups=rts_bwd_group, replica_group=rg_dp, use_io_tiles=USE_IO_TILES ) def fwd_load(self, i: Union[int, popxl.Tensor]): return NamedTensors.pack(self._fwd_load_names, self._fwd_load.call(i)) def grad_store(self, args: NamedTensors, i: Union[float, popxl.Tensor]): return self._grad_store.bind(args).call(i) def optim_fwd_load(self, i: Union[int, popxl.Tensor]): return NamedTensors.pack(self._optim_fwd_load_names, self._optim_fwd_load.call(i)) def optim_fwd_store(self, args: NamedTensors, i: Union[int, popxl.Tensor]): return self._optim_fwd_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 batch_serialise_layer( graphs: Graphs, input_streams: addons.InputStreams, output_streams: addons.OutputStreams, buffers: Dict[str, popxl.RemoteBuffer], shard_group: Optional[popxl.ReplicaGrouping], ): config = graphs.config phase_config: PhaseConf = graphs.layer_config.phase_config shard_groups = { name: shard_group if buffer.meta_shape else gcg().ir.replica_grouping(group_size=1) for name, buffer in buffers.items() } load_handles = {} store_streams = {} store_buffers = {} seed_input = None for io in ("fwd_inputs", "bwd_inputs", "fwd_outputs", "bwd_outputs"): if io == "fwd_inputs": graph_tensors = OrderedDict([(t.name, t) for t in graphs.fwd.graph.inputs]) conf = phase_config.fwd_inputs elif io == "bwd_inputs": graph_tensors = OrderedDict([(t.name, t) for t in graphs.bwd.graph.inputs]) conf = phase_config.bwd_inputs elif io == "fwd_outputs": graph_tensors = OrderedDict([(t.name, t) for t in graphs.fwd.graph.outputs]) conf = phase_config.fwd_outputs elif io == "bwd_outputs": graph_tensors = OrderedDict([(t.name, t) for t in graphs.bwd.graph.outputs]) conf = phase_config.bwd_outputs for name_or_idx, handle in conf.items(): if isinstance(name_or_idx, str): t = graph_tensors[name_or_idx] else: t = graph_tensors.idx(name_or_idx) if handle.type == "stream": if "inputs" in io: stream = input_streams[handle.name] load_handles[t] = stream else: stream = output_streams[handle.name] store_streams[t] = stream elif handle.type == "seed": assert seed_input is None and "inputs" in io seed_input = t elif handle.type == "buffer": buffer = buffers[handle.name] shard_group = shard_groups[handle.name] if handle.rts else None remote = RemoteHandle(buffer, handle.row_offset, shard_group) if "inputs" in io: load_handles[t] = remote else: store_buffers[t] = remote else: raise Exception("unknown type") if not phase_config.fwd_only: fwd, bwd = batch_serialise_fwd_and_grad( graphs.fwd, graphs.bwd, graphs.fwd.args, config.gradient_accumulation, load_handles=load_handles, store_streams=store_streams, store_buffers=store_buffers, seed_input=seed_input, rows=phase_config.rows, io_mode="io", ) graphs.fwd = fwd.graph graphs.bwd = bwd.graph else: fwd = batch_serialise( graphs.fwd, config.gradient_accumulation, load_handles=load_handles, store_streams=store_streams, store_buffers=store_buffers, seed_input=seed_input, rows=phase_config.rows, io_mode="io", ) graphs.fwd = fwd.graph def optimizer_step(optim_graphs: OptimGraphs, ts: NamedTensors, lr: popxl.Tensor, global_norm: popxl.Tensor): _variables = ts.fwd.to_dict() _state = ts.optim _grads = ts.bwd.accum.to_dict() for name, graph in optim_graphs.items(): graph.bind(_state[name]).call(_variables[name], _grads[name], lr, global_norm) def task_head_optimizer_step(optim_graphs: OptimGraphs, ts: NamedTensors, lr: popxl.Tensor, global_norm: popxl.Tensor): _variables = ts.fwd.to_dict() _state = ts.optim _grads = {name.replace("accum.", ""): t for name, t in ts.bwd.to_dict().items()} for name, graph in optim_graphs.items(): graph.bind(_state.get(name)).call(_variables[name], _grads[name], lr, global_norm) def global_norm_reduce(config: GPTConfig, grad_norm: popxl.Tensor, grads: NamedTensors): for g in grads.tensors: ops.add_(grad_norm, grad_reduce_square_add(g, config.execution.loss_scaling)) def pretraining(config: GPTConfig) -> TaskSession: replicas = config.execution.data_parallel * config.execution.tensor_parallel ir = popxl.Ir(replication="popdist" if popdist.isPopdistEnvSet() else replicas) assert ir.replication_factor == replicas layer_config = gen_layer_config(config) # Options opts = ir._pb_ir.getSessionOptions() opts.numIOTiles = config.execution.io_tiles opts.enableStochasticRounding = config.training.stochastic_rounding opts.partialsTypeMatMuls = "half" opts.engineOptions["target.syncReplicasIndependently"] = "true" main = ir.main_graph with timer("PopXL IR construction"): with main: rg_tp, rg_dp = replica_groups(config) rg_rts_activations = rg_tp # ----- Define input and output streams ----- input_shape = (config.execution.micro_batch_size * config.model.sequence_length,) input_streams = addons.InputStreams( words=(input_shape, popxl.int32), labels=(input_shape, popxl.int32), lr=((), popxl.float32) ) output_streams = addons.OutputStreams(loss=((), config.model.dtype), grad_norm=((), popxl.float32)) positions = popxl.constant(generate_positions(config), popxl.int32, name="positions") # ---- Initialise Random Seed ---- # Same seed for tp1 group. Different across tp2+dp groups seed_v, seed = addons.seed_variable(config.model.seed, replica_grouping=rg_tp) # ----- Build compute graphs ----- optimizer = AdamOptimizerStep() embeddings = Graphs( config, layer_config, optimizer, GPTEmbeddingsTP, 1, None, input_streams.words.spec, positions.spec, seed=seed.spec, ) x_spec = embeddings.fwd.graph.outputs[0] decoder_block = Graphs( config, layer_config, optimizer, GPTDecoderBlockTP, config.model.layers, None, x_spec, seed=seed.spec ) tied_weight_spec = embeddings.fwd.args.word.weight head = Graphs( config, layer_config, optimizer, GPTLMHeadLossTP, 1, {"lm_head.word_embedding": embeddings.buffers.fwd.word.weight}, x_spec, input_streams.labels.spec, ) # Make Head a single Fwd+Bwd layer to improve phase efficiency _, head.fwd = HeadFwdBwdTiedEmb(config, head.fwd, head.bwd, head.facts.fwd, head.grad_facts).create_graph( x_spec, input_streams.labels.spec, tied_weight_spec, tied_weight_spec ) # ---- Transform graphs ---- # Recomputation embeddings.bwd = addons.recompute_graph(embeddings.bwd) decoder_block.bwd = addons.recompute_graph(decoder_block.bwd) # Batch Serialisation # Buffers act_shard_group = ( get_activ_shard_group(x_spec, rg_rts_activations, RTS_ACTIVATIONS_THRESHOLD) if RTS_ACT else None ) x_buffer = batch_serial_buffer( embeddings.fwd.graph.outputs[0], steps=config.gradient_accumulation, rows=config.model.layers + 1, shard_group=act_shard_group, ) dx_buffer = batch_serial_buffer( embeddings.bwd.graph.inputs[0], steps=config.gradient_accumulation, rows=config.model.layers + 1, shard_group=act_shard_group, ) buffers = {"x": x_buffer, "dx": dx_buffer} # Graphs batch_serialise_layer(embeddings, input_streams, output_streams, buffers, act_shard_group) batch_serialise_layer(decoder_block, input_streams, output_streams, buffers, act_shard_group) batch_serialise_layer(head, input_streams, output_streams, buffers, act_shard_group) # Available Memory Proportion addons.set_available_memory_proportion_by_ipu(ir, config.execution.available_memory_proportion) # ----- Create Variables ----- # Structure should match gpt_lm.GPTLMHeadModelTP.hf_mapping variables = NamedTensors(random_seed=seed_v) transformer = NamedTensors() variables.insert("transformer", transformer) transformer.insert( "embeddings", embeddings.remote_buffer_facts.init_remote( embeddings.buffers, 0, "embeddings", ), ) transformer.insert( "decoder", NamedTensors.from_dict( { n: decoder_block.facts.init_remote( decoder_block.buffers, n, f"decoder.{n}", ) for n in range(config.model.layers) } ), ) variables.insert( "head", head.facts.init_remote( head.buffers, 0, "head", ), ) # ---- Execute ---- with popxl.in_sequence(): # Load current learning rate lr = ops.host_load(input_streams.lr) # Increment random seed seed += 1 def embedding_fwd_phase(seed: popxl.Tensor, positions: popxl.Tensor): # 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, positions) return seed embed_fwd_graph = ir.create_graph(embedding_fwd_phase, seed, positions) (seed,) = ops.call(embed_fwd_graph, seed, positions) def single_decoder_block_fwd_phase(n: popxl.Tensor, seed: popxl.Tensor): # Load decoder block layer_vars = decoder_block.fwd_load(n) layer_vars = decoder_block.fwd_all_gather(layer_vars) # Forward seed, layer_seed = ops.split_random_seed(seed) decoder_block.fwd.bind(layer_vars).call(n, layer_seed) return n + 1, seed i = popxl.constant(0, name="layer_index") fwd_graph = ir.create_graph(single_decoder_block_fwd_phase, i, seed) ops.repeat(fwd_graph, config.model.layers, i, seed) # Buffer to be used later tied_weight_grad_buffer = None def task_head_fwd_grad_phase(): nonlocal tied_weight_grad_buffer # Load task head layer head_vars = NamedTensors(fwd=head.fwd_all_gather(head.fwd_load(0)), bwd=head.grad_facts.init_zero()) # Tied weight tied_weight = head_vars.fwd.lm_head.pop("word_embedding") tied_weight_grad = ops.init(tied_weight.shape, tied_weight.dtype, "word_embedding_grad", "zero") # Forward + Gradient head.fwd.bind(head_vars).call(0, tied_weight, tied_weight_grad) # Data parallel reduce reduced_grads = head.grad_reduce(head_vars.bwd) # Global Norm calculation grad_norm = ops.init((), popxl.float32, name="grad_norm", init_type="zero") global_norm_reduce(config, grad_norm, reduced_grads) # Store Gradients head.grad_store(reduced_grads, 0) # Reduce and Store the tied gradient grad_t = reduce_replica_sharded_tensor( tied_weight_grad, "mean", replica_group=rg_dp, shard_group=rg_dp ) tied_weight_grad_buffer = create_remote_buffer( grad_t.spec, replica_group=rg_dp, shard_over=rg_dp.group_size ) ops.remote_store(tied_weight_grad_buffer, 0, grad_t) return grad_norm task_graph = ir.create_graph(task_head_fwd_grad_phase) (grad_norm,) = ops.call(task_graph) def single_decoder_block_grad_phase(n: popxl.Tensor, grad_norm: popxl.TensorByRef): # Load layer layer_vars = decoder_block.fwd_load(n) layer_vars = decoder_block.fwd_all_gather(layer_vars) # Gradient grads = decoder_block.grad_facts.init_zero() bwd_vars = grads.copy() bwd_vars.update(layer_vars) decoder_block.bwd.bind(bwd_vars).call(n) # Data parallel reduce reduced_grads = decoder_block.grad_reduce(grads) # Global Norm calculation global_norm_reduce(config, grad_norm, reduced_grads) # Store gradient decoder_block.grad_store(reduced_grads, n) return n - 1 i = popxl.constant(config.model.layers - 1, name="layer_index") bwd_graph = ir.create_graph(single_decoder_block_grad_phase, i, grad_norm) ops.repeat(bwd_graph, config.model.layers, i, grad_norm) def embedding_grad_optimizer_phase(lr: popxl.Tensor, grad_norm: popxl.TensorByRef): nonlocal tied_weight_grad_buffer # Load Embeddings layer embeddings_vars = embeddings.optim_fwd_load(0) embeddings_fwd_vars = embeddings.fwd_all_gather(embeddings_vars.fwd) # Gradient grads = embeddings.grad_facts.init_zero() bwd_vars = grads.copy() bwd_vars.update(embeddings_fwd_vars) embeddings.bwd.bind(bwd_vars).call(0) # Data parallel reduce reduced_grads = embeddings.grad_reduce(grads) # Add the tied gradient from the projection tied_weight_grad = ops.remote_load(tied_weight_grad_buffer, 0) ops.add_(reduced_grads.accum.word.weight, tied_weight_grad) # Global Norm calculation global_norm_reduce(config, grad_norm, reduced_grads) # Finalise global bwd norm with an all reduce and sqrt grad_norm = ops.sqrt(ops.collectives.replicated_all_reduce(grad_norm, op="add")) ops.host_store(output_streams.grad_norm, grad_norm) # Optimizer Step for Embeddings. # Note: No need to store then load the gradient.. just use it directly embeddings_vars.insert("bwd", reduced_grads) optimizer_step(embeddings.optim, embeddings_vars, lr, grad_norm) # Store embeddings.optim_fwd_store(embeddings_vars, 0) return grad_norm embed_bwd_graph = ir.create_graph(embedding_grad_optimizer_phase, lr, grad_norm) (grad_norm,) = ops.call(embed_bwd_graph, lr, grad_norm) # Optimizer Step for Layers def layer_optim(n: popxl.Tensor, lr: popxl.Tensor, grad_norm: popxl.Tensor): layer_vars = decoder_block.optim_fwd_load(n) optimizer_step(decoder_block.optim, layer_vars, lr, grad_norm) decoder_block.optim_fwd_store(layer_vars, n) return n + 1 i = popxl.constant(0, name="layer_index") layer_optim_graph = ir.create_graph(layer_optim, i, lr, grad_norm) ops.repeat(layer_optim_graph, config.model.layers, i, lr, grad_norm) def head_optim(lr: popxl.Tensor, grad_norm: popxl.Tensor): # Optimizer Step for Task Head - Only layer norm, tied weights handled by embedding head_vars = head.optim_fwd_load(0) task_head_optimizer_step(head.optim, head_vars, lr, grad_norm) # Store head.optim_fwd_store(head_vars, 0) head_optim_graph = ir.create_graph(head_optim, lr, grad_norm) ops.call(head_optim_graph, lr, grad_norm) # Run `OpToIdentityPattern` among others part of `PreAliasPatterns` apply_pre_alias_patterns(ir, level="default") repeat_graph(main, config.execution.device_iterations) fwd_vars = NamedTensors.from_dict( { "transformer.embeddings": variables.transformer.embeddings.fwd, "transformer.decoder": NamedTensors.from_dict( {i: variables.transformer.decoder[i].fwd for i in range(config.model.layers)} ), "head": variables.head.fwd, } ) ir.num_host_transfers = config.execution.device_iterations * config.gradient_accumulation session = TaskSession( input_streams, output_streams, fwd_vars, ir=ir, device_desc="ipu_hw", ) return session def main(): """Run a benchmark configuration""" config, _, _ = gpt_config_setup( CONFIG_DIR / "pretraining.yml", "release", "tiny", wandb_setup=False, hf_model_setup=False ) session = pretraining(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 i in range(5): start = time.perf_counter() outputs = session.run(inputs) loss = outputs[session.outputs[0]].mean() durations.append(time.perf_counter() - start) logging.info(f"Step {i}. Loss {loss:.2f}") 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/s " logging.info(result_str) if __name__ == "__main__": try: main() except Exception as e: logging.exception(e, exc_info=False) # Log time of exception raise