File size: 2,369 Bytes
26d5b81 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | #ifndef NEUROFLOW_DPO_HPP
#define NEUROFLOW_DPO_HPP
#include <cstddef>
#include <random>
#include <string>
#include <vector>
#include "adamw.hpp"
#include "alignment_common.hpp"
#include "causal_lm.hpp"
#include "scheduler.hpp"
#include "tokenizer.hpp"
namespace neuroflow {
struct DPOTrainConfig {
std::string data_path;
std::string sft_ckpt_path;
std::string tokenizer_path;
std::string output_dir;
float learning_rate = 1e-6f;
int epochs = 3;
size_t max_seq_len = 512;
float warmup_ratio = 0.05f;
float weight_decay = 0.0f;
float grad_clip = 1.0f;
float beta = 0.1f;
float adam_beta1 = 0.9f;
float adam_beta2 = 0.999f;
float adam_eps = 1e-8f;
size_t save_interval = 1000;
size_t log_interval = 10;
unsigned seed = 42;
};
class DPODataLoader {
public:
DPODataLoader(const std::string& jsonl_path, size_t max_samples = 0);
bool has_next() const;
DPOSample next();
void reset();
void shuffle(std::mt19937& rng);
size_t total_samples() const { return samples_.size(); }
size_t invalid_count() const { return invalid_count_; }
private:
std::vector<DPOSample> samples_;
size_t cursor_ = 0;
size_t invalid_count_ = 0;
};
struct DPOLossOutput {
float loss;
float alpha;
float reward_chosen;
float reward_rejected;
};
float compute_log_prob(CausalLMHead& model, const std::vector<size_t>& token_ids,
size_t prompt_len, size_t vocab_size);
DPOLossOutput compute_dpo_loss(float log_prob_chosen_policy,
float log_prob_rejected_policy,
float log_prob_chosen_ref,
float log_prob_rejected_ref,
float beta);
class DPOTrainer {
public:
DPOTrainConfig config;
DPOTrainer(const DPOTrainConfig& cfg);
void train();
float train_on_sample(const DPOSample& sample);
private:
std::unique_ptr<CausalLMHead> policy_;
std::unique_ptr<CausalLMHead> reference_;
std::unique_ptr<BPETokenizer> tokenizer_;
std::unique_ptr<AdamW> optimizer_;
std::unique_ptr<CosineScheduler> scheduler_;
float compute_w_embed_checksum();
};
} // namespace neuroflow
#endif // NEUROFLOW_DPO_HPP |