)jgddlmZddlmZdZdZdZeZeZ ej dZ e j dZ dZ dZd ZeZeZej d Zej d Zd ZdS) NctjtjkotjS)N)mxdefault_devicegpumetal is_available]/lsinfo/ai/hellotax_ai/base_platform/venv/lib/python3.11/site-packages/mlx_lm/tuner/losses.py can_run_metalr s+   "& ( DRX-B-B-D-DDr cttsdSd}tjdddgdg|dS)Naz constexpr int M = 4; constexpr int block = 1024 * M; constexpr int full_blocks = V / block; constexpr int extra = V - full_blocks * block; threadgroup float shared[32 * 2]; uint out_idx = threadgroup_position_in_grid.y; uint simd_lane_id = thread_index_in_simdgroup; uint simd_group_id = simdgroup_index_in_threadgroup; logits_q += out_idx * V; logits_p += out_idx * V; out += out_idx; float lse_q_minus_p; float lse_p; { float max_q = -1e30; float max_p = -1e30; float sum_exp_q = 0; float sum_exp_p = 0; int offset = thread_index_in_threadgroup * M; for (int i = 0; i < full_blocks; i++) { // Read and update q and p float vals_q[M]; float vals_p[M]; for (int j=0; j 0) { // Read and update q and p float vals_q[M]; float vals_p[M]; for (int j=0; j < M; j++) { vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30; vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30; } float prev_max_q = max_q; float prev_max_p = max_p; for (int j=0; j 0) { float vals_q[M]; float vals_p[M]; for (int j=0; j(kl); } } kl_forwardlogits_qlogits_poutTname input_names output_namessourceensure_row_contiguousr rfast metal_kernelrs r _make_kl_forward_kernelr sS ??ZFx 7   ,W"   r cttsdSd}tjdgddg|dS)Na constexpr int M = 4; constexpr int block = 1024 * M; constexpr int full_blocks = V / block; constexpr int extra = V - full_blocks * block; threadgroup float shared[32 * 2]; uint out_idx = threadgroup_position_in_grid.y; uint simd_lane_id = thread_index_in_simdgroup; uint simd_group_id = simdgroup_index_in_threadgroup; logits_q += out_idx * V; logits_p += out_idx * V; out += out_idx * V; cotan += out_idx; float lse_q; float lse_p; { float max_q = -1e30; float max_p = -1e30; float sum_exp_q = 0; float sum_exp_p = 0; int offset = thread_index_in_threadgroup * M; for (int i = 0; i < full_blocks; i++) { // Read and update q and p float vals_q[M]; float vals_p[M]; for (int j=0; j 0) { // Read and update q and p float vals_q[M]; float vals_p[M]; for (int j=0; j < M; j++) { vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30; vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30; } float prev_max_q = max_q; float prev_max_p = max_p; for (int j=0; j( c * (metal::fast::exp(vals_q[j] - lse_q) - metal::fast::exp(vals_p[j] - lse_p))); } // Move to the next block offset += block; } if (extra > 0) { float vals_q[M]; float vals_p[M]; for (int j=0; j( c * (metal::fast::exp(vals_q[j] - lse_q) - metal::fast::exp(vals_p[j] - lse_p))); } } } } kl_backward)rrcotanrTrrrs r _make_kl_backward_kernelr sT ??SFj 7   555W"   r c |j|jdz}|j}t||g|jddg|gd|fd|jdfgd|dfddS NTVr&r'r'inputs output_shapes output_dtypestemplategrid threadgroupr)sizeshapedtype_kl_forward_kernel)rrn_outsdts r _kl_div_lossr6Xs ]hnR0 0F B (#~crc*+d)c8>"#567FA       r c |\}}|j}tj|}t|||g|jg|gd|fd|jdfgd|jdfdd}||fS) Nr$r%r#r&r'r(r)r)r2r zeros_like_kl_backward_kernelr1r0)primals cotangentoutputrrr5dpdqs r r6r6gs Hh B x B (I.~&d)c8>"#567INA &     B r6Mr c trt||Stj|t j|ddz |t j|ddz ddS)Nr#TaxiskeepdimsnonerA reduction)r r6nnlosses kl_div_lossr logsumexp)rrs r rHrHysu Hh///y$$ r|H2EEE E r|H2EEE E %   r cvtsdSd}tjdddgddg|dS) NaT constexpr int M = 4; constexpr int block = 1024 * M; constexpr int full_blocks = V / block; constexpr int extra = V - full_blocks * block; threadgroup float shared[32 * 2]; uint out_idx = threadgroup_position_in_grid.y; uint simd_lane_id = thread_index_in_simdgroup; uint simd_group_id = simdgroup_index_in_threadgroup; logits_q += out_idx * V; logits_p += out_idx * V; out += out_idx; out_kl_q += out_idx; float lse_p; float lse_q; { float max_q = -1e30; float max_p = -1e30; float sum_exp_q = 0; float sum_exp_p = 0; int offset = thread_index_in_threadgroup * M; for (int i = 0; i < full_blocks; i++) { // Read and update q and p float vals_q[M]; float vals_p[M]; for (int j=0; j 0) { // Read and update q and p float vals_q[M]; float vals_p[M]; for (int j=0; j < M; j++) { vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30; vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30; } float prev_max_q = max_q; float prev_max_p = max_p; for (int j=0; j(2)); int offset = thread_index_in_threadgroup * M; for (int i = 0; i < full_blocks; i++) { // Read and add to the kl_p and kl_q float vals_q[M]; float vals_p[M]; for (int j=0; j 0) { float vals_q[M]; float vals_p[M]; for (int j=0; j(0.5 * kl_p + 0.5 * kl_q); out_kl_q[0] = static_cast(kl_q); } } js_forwardrrrout_kl_qTrrrs r _make_js_forward_kernelrMsV ??lF\ 7   ,Z("   r cttsdSd}tjdgddg|dS)Na constexpr int M = 4; constexpr int block = 1024 * M; constexpr int full_blocks = V / block; constexpr int extra = V - full_blocks * block; threadgroup float shared[32 * 2]; uint out_idx = threadgroup_position_in_grid.y; uint simd_lane_id = thread_index_in_simdgroup; uint simd_group_id = simdgroup_index_in_threadgroup; logits_q += out_idx * V; logits_p += out_idx * V; out_q += out_idx * V; cotan += out_idx; output_kl_q += out_idx; float lse_q; float lse_p; { float max_q = -1e30; float max_p = -1e30; float sum_exp_q = 0; float sum_exp_p = 0; int offset = thread_index_in_threadgroup * M; for (int i = 0; i < full_blocks; i++) { // Read and update q and p float vals_q[M]; float vals_p[M]; for (int j=0; j 0) { // Read and update q and p float vals_q[M]; float vals_p[M]; for (int j=0; j < M; j++) { vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30; vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30; } float prev_max_q = max_q; float prev_max_p = max_p; for (int j=0; j(2)); float kl_q = output_kl_q[0]; int offset = thread_index_in_threadgroup * M; for (int i = 0; i < full_blocks; i++) { // Read and compute vjp for logits_q float vals_q[M]; float vals_p[M]; for (int j=0; j( c * 0.5 * q_j * (logtwo - metal::fast::log(1 + metal::fast::exp(logp_j - logq_j)) - kl_q) ); } // Move to the next block offset += block; } if (extra > 0) { float vals_q[M]; float vals_p[M]; for (int j=0; j( c * 0.5 * q_j * (logtwo - metal::fast::log(1 + metal::fast::exp(logp_j - logq_j)) - kl_q) ); } } } } js_backward)rrr output_kl_qout_qTrrrs r _make_js_backward_kernelrR?sT ??]F~ 7   DDDY"   r c |j|jdz}|j}t||g|jdd|jddg||gd|fd|jdfgd|dfd}|dt j|dfSr")r0r1r2_js_forward_kernelr stop_gradient)rrr4r5outputss r _js_div_lossrWs ]hnR0 0F B (#~crc*HN3B3,?@2h)c8>"#567FA  G 1:r' 33 33r c |\}}|\}}|\}}|j}tj|} t||||g|jg|gd|fd|jdfgd|jdfd} | | fS)Nr$r%r#r&r'r(r))r2rr8_js_backward_kernelr1r0) r: cotangentsrVrrr_kl_qr5r=r>s r rWrWs HhHE1GAt B x B (E40~&d)c8>"#567EJ "    B r6Mr ctrt||dS|tj|ddz }|tj|ddz }|tjdtj||z zztjd|jz }tj ||dd}tj ||dd}d ||zzS) Nrr#Tr@r'rCrDg?) r rWrrIlogexpastyper2rFrGrH)rr logprobs_p logprobs_q logprobs_mkl_pr\s r js_div_lossrfs #Hh//22 XB N N NN  XB N N NN fQ Z 788899 :fQiix~.. /  y$$Z"PV$WWy$$Z"PV$WWdTk""r )mlx.corecorermlx.nnrFr rr r3r9custom_functionr6vjprHrMrRrTrYrWrfr r r rlsmEEEeeeP^^^B-,....00     "    wwwthhhV-,....00 4 4 4$ # # # # #r