class CompositePEFT:
    def __init__(self, *pefts):
        self.pefts = pefts

    def parameters(self):
        for p in self.pefts:
            yield from p.parameters()

    def state_dict(self):
        sd = {}
        for i, p in enumerate(self.pefts):
            if not hasattr(p, "state_dict"):
                continue
            p_sd = p.state_dict()
            
            for k, v in p_sd.items():
                if isinstance(v, dict):
                    for k2, v2 in v.items():
                        sd[f"{i}.{k}.{k2}"] = v2
                else:
                    sd[f"{i}.{k}"] = v
        return sd

    def load_state_dict(self, sd):
       
        per = {}
        for key, val in sd.items():
            
            if "." not in key:
                continue
            idx_str, rest = key.split(".", 1)
            idx = int(idx_str)
            per.setdefault(idx, {})[rest] = val
        
        for i, p in enumerate(self.pefts):
            if i not in per:
                continue
            p_sd = per[i]
            if hasattr(p, "load_state_dict"):
                
                p.load_state_dict(p_sd)