def save_exported_model()

in summarize_from_feedback/query_response_model.py [0:0]

• 27 lines of code
• 9 McCabe index (conditional complexity)

def save_exported_model(layout, model, model_H: Hyperparams, save_dir, save_heads: Set[str]):
    """
    Exporting a model allows it to be run with a different layout than it was trained with.

    Currently, uploading/loading an exported model is slower than saving/restoring a checkpoint,
    but if we can get exporting to be sufficiently fast, then we could replace legacy_checkpoints.py with
    this "exporting" approach.
    """
    if blobs.is_blob_url(save_dir):
        local_dir = os.path.join("/tmp", str(uuid.uuid4()))
    else:
        local_dir = save_dir
    os.makedirs(os.path.join(local_dir, "checkpoint"), exist_ok=True)

    def export_fine_piece(fine_model_piece_dict: dict, chkpt_prefix: str):
        fine_piece_path = os.path.join(
            local_dir, "checkpoint", f"{chkpt_prefix}_shard_{layout.shard_idx:03d}.pkl"
        )
        # print(f"Uploading fine_piece: {fine_piece_path}")
        torch.save(fine_model_piece_dict, fine_piece_path)
        torch.cuda.synchronize()  # Verify that the piece has finished being written

    # Export the embeddings
    if model.include_input_embeddings:
        export_fine_piece(model.embedding.state_dict(), "input_embeddings")

        if model.include_pos_embeddings:
            export_fine_piece(model.position_embedding.state_dict(), "position_embedding")
    # Export the resblocks
    for resblock_idx, resblock in enumerate(model.torso.resblocks):
        export_fine_piece(resblock.state_dict(), f"resblock_{resblock_idx:04d}")
    # Export the final_layer_norm
    if model.include_final_layer_norm:
        export_fine_piece(model.ln_f.state_dict(), "final_layer_norm")
    # Export the unembeddings
    if model.include_output_unembeddings:
        export_fine_piece({"unembedding_weights": model.unembedding_weights}, "output_unembeddings")
    for head in save_heads:
        export_fine_piece(model.scalar_heads[head].state_dict(), f"output_head_{head}")
    if blobs.is_blob_url(save_dir):
        blobs.parallel_copy_recursive(local_dir, save_dir)
        shutil.rmtree(local_dir)