I need some help merging two checkpoint.pth files, does my code actually merge well?

username: not_Narx
cooked: <p>Hi, I’m relatively new to PyTorch. I have a Federated AI Training system, so each user trains the model locally from a checkpoint on a server and uploads the modified checkpoint to the server. In my server code, I merge these two by averaging the weights. Here is the code snippet:</p>
<p>def merge_checkpoint_files(existing_path, new_path):</p>
<pre><code class="lang-auto">existing_ckpt = torch.load(existing_path, map_location="cpu")
new_ckpt      = torch.load(new_path,      map_location="cpu")

existing_epoch = existing_ckpt.get("epoch", 0)
new_epoch      = new_ckpt.get("epoch", 0)
existing_state = existing_ckpt.get("model_state_dict", existing_ckpt)
new_state      = new_ckpt.get("model_state_dict", new_ckpt)

# average any matching weights:
merged_state = {}
for key in existing_state.keys():
    if key in new_state and existing_state[key].shape == new_state[key].shape:
        merged_state[key] = (existing_state[key] + new_state[key]) / 2
    else:
        merged_state[key] = existing_state[key]

# copyany new weights which appear in new_state
for key in new_state.keys():
    if key not in merged_state:
        merged_state[key] = new_state[key]

# build new checkpoint.pth so that it keeps all the original keys,
# but overwrites model_state_dict and updates "epoch".
if "model_state_dict" in existing_ckpt:
    merged_ckpt = existing_ckpt.copy()
    merged_ckpt["model_state_dict"] = merged_state
else:
    merged_ckpt = merged_state

# overwrite the "epoch" field
merged_ckpt["epoch"] = max(existing_epoch, new_epoch)

# save
torch.save(merged_ckpt, existing_path)
</code></pre>
<hr>
<p>Does this code merge the files well? In my use case, when a user uploads and the server merges it, after a few seconds, other users may also upload theirs, and will there be confusion or loss of data? Is there a better way? I have tried using the HuggingFace Transformers library, but it doesn’t work for me as I am using Python 3.12</p>
<p>Any answer helps, Thank You for your time <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <p>It’s unclear to me if any why averaging two different <code>state_dict</code>s` should work, i.e. should either improve or keep the same model performance.</p>
<p>Besides that, I believe your code looks fine.</p>

username: not_Narx
cooked: <p>Is there like any other way to merge two of the files together?<br>
Thank you for the reply!</p>

username: ptrblck
cooked: <p>I’m not familiar with any proven approaches and we discussed the approach to average parameters e.g. in <a href="https://discuss.pytorch.org/t/creating-a-model-which-weights-are-the-sum-of-weights-of-2-different-neural-networks/119491/5">this thread</a>. Let us know if you find any references to verified approaches.</p>

username: not_Narx
cooked: <p>Yep, this is a pretty new subject. There aren’t any references to verified approaches that have been proven to work. But there is a Package called FLPytorch which helps in doeing federated AI trainings, but then I have tried but the packages were really buggy for me and some stuff just didnt work. But they also use a similar kind of logic of averaging. The package is also not there in Pycharm, you should download it from Github. When I checked the thread, it’s pretty much the same confusion that I have. Two models trained on similar data with the same configuration may have different weights for a specific token of a word, so when you try to merge them, averaging them might not always work.  However, in my case, since I use the same code to train the model, it may work. I am currently testing it on this merging logic, which in my opinion, is a better approach.</p>
<p>Heres the new merging code:</p>
<p>def merge_checkpoint_files(existing_path, new_path,<br>
vocab_path_existing=VOCAB_FILE,<br>
vocab_path_new=None):</p>
<pre><code>print("Started merging checkpoints.")
ckpt1 = torch.load(existing_path, map_location="cpu")
ckpt2 = torch.load(new_path,     map_location="cpu")
ep1, ep2 = ckpt1.get("epoch",0), ckpt2.get("epoch",0)
s1 = ckpt1["model_state_dict"]
s2 = ckpt2["model_state_dict"]

old_vocab, new_vocab, merged_vocab = merge_vocab_files(
    vocab_path_existing, vocab_path_new or vocab_path_existing
)
rev1 = {i:t for t,i in old_vocab.items()}
rev2 = {i:t for t,i in new_vocab.items()}


merged_state = {}
w1, w2 = ep1+1, ep2+1
for key in set(s1) | set(s2):
    t1, t2 = s1.get(key), s2.get(key)
    if t1 is not None and t2 is not None and t1.shape == t2.shape:
        kl = key.lower()
        if "wte.weight" in kl or "embedding" in kl:
            D = t1.size(1)
            M = len(merged_vocab)
            merged_emb = torch.zeros((M, D), dtype=t1.dtype)
            for token, idx in merged_vocab.items():
                i1 = old_vocab.get(token)
                i2 = new_vocab.get(token)
                if i1 is not None and i2 is not None:
                    merged_emb[idx] = (t1[i1]*w1 + t2[i2]*w2)/(w1+w2)
                elif i1 is not None:
                    merged_emb[idx] = t1[i1]
                else:
                    merged_emb[idx] = t2[i2]
            merged_state[key] = merged_emb

        elif any(tag in kl for tag in ("bias","norm","bn","ln")):
            merged_state[key] = t2.clone()
        else:
            merged_state[key] = (t1*w1 + t2*w2)/(w1+w2)

    elif t1 is not None:
        merged_state[key] = t1.clone()
    else:
        merged_state[key] = t2.clone()

new_ckpt = copy.deepcopy(ckpt1)
new_ckpt["model_state_dict"] = merged_state
new_ckpt["epoch"] = max(ep1, ep2)
torch.save(new_ckpt, existing_path)

with open(vocab_path_existing, "wb") as f:
    pickle.dump(merged_vocab, f)
</code></pre>
<p>In my understanding, when you average the weights with a better-trained model trained in a similar pathway, it might work. But in my code, not everything is averaged tho. Also, I think after merging it, if I train it once again without any merging, it should train over any weights that didn’t really work after the merge.</p>
<p>Anyways, I am currently testing this logic with 300k entries with about 20 epochs using a single computer, which should allow the model to have a basic, I mean BASIC, and decent, small sentence formations. This way I can see if even models trained in the same pathway will work after merging. But I am training it on my computer so it will take a few days, but I will post the results.</p>
<p>Thank you for your time!</p>