diff options
Diffstat (limited to 'experiments/transformer_crossover_native.py')
| -rw-r--r-- | experiments/transformer_crossover_native.py | 92 |
1 files changed, 88 insertions, 4 deletions
diff --git a/experiments/transformer_crossover_native.py b/experiments/transformer_crossover_native.py index 43d3e62..7ccb8fb 100644 --- a/experiments/transformer_crossover_native.py +++ b/experiments/transformer_crossover_native.py @@ -141,10 +141,11 @@ def train_step(net, optimizer, tokens, targets, args, step, generators): metrics = {"loss": float(loss.detach())} elif args.method == "dfa": with torch.no_grad(): - loss = net(tokens, targets)["loss"] + cache = net(tokens, targets, return_cache=True) + loss = cache["loss"] metrics = { "loss": float(loss), - **net.dfa_gradients(tokens, targets), + **net.dfa_gradients(tokens, targets, cache=cache), } elif args.method == "pepita": local = net.pepita_gradients(tokens, targets) @@ -270,7 +271,9 @@ def hardware_report(device): return report -def work_report(net, args, train_tokens, validation_evaluations): +def work_report( + net, args, train_tokens, completed_steps, + validation_evaluations): ordinary_validation = sum( row["tokens"] for row in validation_evaluations) presentations = train_tokens @@ -291,6 +294,70 @@ def work_report(net, args, train_tokens, validation_evaluations): train_tokens * (args.ep_free_steps + args.ep_nudge_steps) + sum(row["relaxation_token_passes"] for row in validation_evaluations)) + width = args.width + block_affine_macs_per_token = ( + (4 + 2 * args.mlp_ratio) * width * width) + block_attention_macs_per_token = ( + 2 * args.context_length * width) + forward_affine_macs_per_token = ( + args.depth * block_affine_macs_per_token + + width * net.config.vocab_size) + forward_attention_macs_per_token = ( + args.depth * block_attention_macs_per_token) + training_forward_multiplier = { + "bp": 1, + "fa": 1, + "dfa": 1, + "pepita": 2, + "ff": 2, + "ep": 2, + "dualprop": 1, + "clean_kp": 1, + "sdil": 1, + }[args.method] + if args.method == "ff": + validation_forward_tokens = sum( + row["candidate_token_presentations"] + for row in validation_evaluations) + else: + validation_forward_tokens = ordinary_validation + full_forward_token_passes = ( + training_forward_multiplier * train_tokens + + validation_forward_tokens) + + local_block_token_evaluations = 0 + local_head_token_evaluations = 0 + if args.method in {"dfa", "pepita"}: + local_block_token_evaluations = args.depth * train_tokens + local_head_token_evaluations = train_tokens + elif args.method == "ff": + steps_per_layer = math.ceil( + args.train_steps / net.ff_num_layers) + block_steps = 0 + head_steps = 0 + for step in range(completed_steps): + layer = min( + step // steps_per_layer, net.ff_num_layers - 1) + block_steps += int(1 <= layer <= args.depth) + head_steps += int(layer == args.depth + 1) + tokens_per_step = args.batch_size * args.context_length + local_block_token_evaluations = block_steps * tokens_per_step + local_head_token_evaluations = head_steps * tokens_per_step + elif args.method == "dualprop": + multiplier = 2 * args.dp_inference_passes + 1 + local_block_token_evaluations = ( + args.depth * train_tokens * multiplier) + local_head_token_evaluations = train_tokens * multiplier + elif args.method == "ep": + train_multiplier = ( + 2 * (args.ep_free_steps + args.ep_nudge_steps) + 2) + validation_multiplier = 2 * args.ep_free_steps + local_block_token_evaluations = args.depth * ( + train_tokens * train_multiplier + + ordinary_validation * validation_multiplier) + local_head_token_evaluations = ( + train_tokens * train_multiplier + + ordinary_validation * validation_multiplier) return { "ordinary_training_tokens": train_tokens, "ordinary_validation_tokens": ordinary_validation, @@ -300,6 +367,22 @@ def work_report(net, args, train_tokens, validation_evaluations): for row in validation_evaluations), "relaxation_token_passes": relaxation, "local_vjp_token_evaluations": local_vjps, + "full_forward_token_passes": full_forward_token_passes, + "local_block_token_evaluations": + local_block_token_evaluations, + "local_head_token_evaluations": local_head_token_evaluations, + "block_affine_macs_per_token": + block_affine_macs_per_token, + "block_attention_macs_per_token": + block_attention_macs_per_token, + "forward_affine_macs_per_token": + forward_affine_macs_per_token, + "forward_attention_macs_per_token": + forward_attention_macs_per_token, + "enumerated_full_forward_macs": ( + full_forward_token_passes + * (forward_affine_macs_per_token + + forward_attention_macs_per_token)), "forward_parameter_count": net.n_forward_parameters, "feedback_parameter_count": net.n_feedback_parameters, "completed_optimizer_steps": len( @@ -434,7 +517,8 @@ def run(args): }, "final": final, "work": work_report( - net, args, train_tokens, validation_evaluations), + net, args, train_tokens, final["step"], + validation_evaluations), "hardware": { **hardware_report(args.device), "peak_memory_allocated_bytes": peak_allocated, |