Rewrite bash scripts into Python interfaces

[rinapch]

Changelog:

• Rewrite 4 main .sh scripts needed for experiment reproduction to Python
• Switch project to poetry
• Remove hardcode from get_dataset function, now you just need to provide the name and the load method

Questions for discussion:

• The main role of the 2 bash scripts is to just set up some variables names and run a Python subprocess with the specified variables. We can just call the Python subprocess directly, however, with a trade-off that we will need to specify a lot of the variables by hand in the CLI
• Should I also convert the training script for LoRA (the get the final model finetuned on selected data)? I didn't use as we have our own training script, but may be useful?

[mishaevtikhiev]

Shouldn't this one be required? Looks like with None the script will fail at the dataset loading stage

[rinapch]

You're right, with None it would fail (fixed it). I made it optional before because for hardcoded initial datasets like tydiqa we don't need to specify it, but realistically we will only run our own datasets, so it's required now

[mishaevtikhiev]

Sanity check question: Adam and AdamW provide the same type of gradient data, am I right?

[rinapch]

In this repo in particular AdamW is used across all computations (despite this variable being just "Adam") -- it's used for warming up the model and for getting the gradients for the train data. For obtaining the gradients of eval data, gradient_type is set to "sgd"

Here is the difference in computation:

• SGD

def obtain_gradients(model, batch):
    """ obtain gradients. """
    loss = model(**batch).loss
    loss.backward()
    vectorized_grads = torch.cat(
        [p.grad.view(-1) for p in model.parameters() if p.grad is not None])
    return vectorized_grads

• Adam(W)

def prepare_optimizer_state(model, optimizer_state, device):
    names = [n for n, p in model.named_parameters() if p.requires_grad]
    avg = torch.cat([optimizer_state[n]["exp_avg"].view(-1) for n in names])
    avg_sq = torch.cat([optimizer_state[n]["exp_avg_sq"].view(-1)
                       for n in names])
    avg = avg.to(device)
    avg_sq = avg_sq.to(device)
    return avg, avg_sq


def obtain_gradients_with_adam(model, batch, avg, avg_sq):
    """ obtain gradients with adam optimizer states. """
    beta1 = 0.9
    beta2 = 0.999
    eps = 1e-08

    loss = model(**batch).loss
    loss.backward()

    vectorized_grads = torch.cat(
        [p.grad.view(-1) for n, p in model.named_parameters() if p.grad is not None])

    updated_avg = beta1 * avg + (1 - beta1) * vectorized_grads
    updated_avg_sq = beta2 * avg_sq + (1 - beta2) * vectorized_grads ** 2
    vectorized_grads = updated_avg / torch.sqrt(updated_avg_sq + eps)

    return vectorized_grads

Looking at this, I think in principle this code would work with a normal Adam, but the gradient values produces by Adam and AdamW would be different

[mishaevtikhiev]

This is the original code from their .sh script, right?

[rinapch]

Yes, it seems like models larger than llama-7B would not fit into their GPUs (relatable) so they needed to enable FSDP for larger experiments